PR #3226: Create optimizers directory and move some files to utils

Imported from GitHub PR #3226

# Description

* Create optimizers directory
* Move some remaining files from `src/MaxText` to `src/maxtext/utils`

# Tests

```
python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
    run_name=<run_name> \
    base_output_directory=gs://<gcs_bucket> \
    dataset_type=synthetic \
    steps=10 \
    enable_checkpointing=false
```

# Checklist

Before submitting this PR, please make sure (put X in square brackets):
- [x] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label.
- [x] I have necessary comments in my code, particularly in hard-to-understand areas.
- [x] I have run end-to-end tests tests and provided workload links above if applicable.
- [x] I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

Copybara import of the project:

--
1e4bdcb by Branden Vandermoon <bvandermoon@google.com>:

Create optimizers directory and move some files to utils

Merging this change closes #3226

COPYBARA_INTEGRATE_REVIEW=#3226 from AI-Hypercomputer:bvandermoon-restructure-last-files 1e4bdcb
PiperOrigin-RevId: 874891637

Author: bvandermoon

src/MaxText/generate_param_only_checkpoint.py

@@ -29,12 +29,12 @@
 import jax
 from jax import random
 from jax.sharding import Mesh
-from MaxText import optimizers
 from MaxText import pyconfig
 from maxtext.common import checkpointing
 from maxtext.common.common_types import DecoderBlockType, MODEL_MODE_TRAIN
 from maxtext.layers import quantizations
 from maxtext.models import models
+from maxtext.optimizers import optimizers
 from maxtext.utils import gcs_utils
 from maxtext.utils import lora_utils
 from maxtext.utils import max_logging

src/MaxText/sft_trainer.py

@@ -28,7 +28,6 @@
 from flax.linen import partitioning as nn_partitioning
 
 from MaxText import pyconfig
-from MaxText import sharding
 from maxtext.trainers.pre_train.train import (
     eval_step,
     get_first_step,
@@ -48,6 +47,7 @@
 from maxtext.utils import max_utils
 from maxtext.utils import max_logging
 from maxtext.utils import maxtext_utils
+from maxtext.utils import sharding
 from maxtext.utils import train_utils
 
 

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

@@ -42,13 +42,13 @@
 import jax
 from jax import random
 from jax.sharding import Mesh
-from MaxText import optimizers
 from MaxText import pyconfig
 from MaxText.globals import MAXTEXT_PKG_DIR
 from maxtext.common import checkpointing
 from maxtext.common.common_types import MODEL_MODE_TRAIN
 from maxtext.layers import quantizations
 from maxtext.models.models import transformer_as_linen
+from maxtext.optimizers import optimizers
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils

src/maxtext/common/data_loader.py

@@ -19,13 +19,13 @@
 import jax.numpy as jnp
 from jax.experimental import checkify
 
-from MaxText.sharding import get_input_data_sharding
 from maxtext.common.goodput import (
     GoodputEvent,
     maybe_record_goodput,
 )
-from maxtext.utils import exceptions
 from maxtext.trainers.diloco import diloco
+from maxtext.utils import exceptions
+from maxtext.utils.sharding import get_input_data_sharding
 
 
 class DataLoader:

src/maxtext/experimental/rl/grpo_trainer.py

@@ -67,7 +67,6 @@
 from ml_goodput_measurement.src.goodput import GoodputRecorder
 
 import MaxText as mt
-from MaxText import sharding
 from MaxText import pyconfig
 from MaxText.globals import EPS
 from maxtext.trainers.pre_train.train import get_first_step
@@ -89,6 +88,7 @@
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils
+from maxtext.utils import sharding
 from maxtext.utils import train_utils
 
 # pylint: disable=too-many-positional-arguments

src/maxtext/inference/paged_attention.py

@@ -30,7 +30,7 @@
 from maxtext.inference import page_manager
 from maxtext.inference import paged_attention_kernel_v2
 from maxtext.layers.initializers import variable_to_logically_partitioned
-from MaxText.sharding import logical_to_mesh_axes
+from maxtext.utils.sharding import logical_to_mesh_axes
 
 _use_kernel_v2 = False
 

src/maxtext/layers/attention_mla.py

@@ -62,7 +62,6 @@
     DEFAULT_MASK_VALUE,
 )
 
-from MaxText.sharding import create_sharding
 from maxtext.layers import nnx_wrappers
 from maxtext.layers.attentions import Attention
 from maxtext.layers.initializers import nd_dense_init, NdInitializer, variable_to_logically_partitioned
@@ -73,6 +72,7 @@
 from maxtext.inference import page_manager
 from maxtext.inference import paged_attention
 from maxtext.inference.kvcache import KVQuant
+from maxtext.utils.sharding import create_sharding
 
 
 class Indexer(nnx.Module):

src/maxtext/layers/attention_op.py

@@ -77,8 +77,8 @@
 from maxtext.layers import nnx_wrappers
 from maxtext.layers.initializers import variable_to_logically_partitioned
 from maxtext.layers.quantizations import AqtQuantization as Quant
-from MaxText.sharding import logical_to_mesh_axes, maybe_shard_with_name
 from maxtext.utils import max_utils
+from maxtext.utils.sharding import logical_to_mesh_axes, maybe_shard_with_name
 import numpy as np
 from tokamax._src.ops.experimental.tpu.splash_attention import splash_attention_kernel as tokamax_splash_kernel
 from tokamax._src.ops.experimental.tpu.splash_attention import splash_attention_mask as tokamax_splash_mask

src/maxtext/layers/attentions.py

@@ -53,7 +53,6 @@
     EP_AS_CONTEXT,
     AttentionType,
 )
-from MaxText.sharding import maybe_shard_with_logical, create_sharding
 from maxtext.layers import nnx_wrappers
 from maxtext.layers.attention_op import AttentionOp
 from maxtext.layers.embeddings import (
@@ -71,6 +70,7 @@
 from maxtext.layers.quantizations import AqtQuantization as Quant
 from maxtext.inference import kvcache, page_manager, paged_attention
 from maxtext.inference.kvcache import KVQuant
+from maxtext.utils.sharding import maybe_shard_with_logical, create_sharding
 
 # pylint: disable=line-too-long, g-doc-args, g-doc-return-or-yield, bad-continuation, g-inconsistent-quotes
 # pytype: disable=attribute-error

src/maxtext/layers/decoders.py

@@ -27,7 +27,6 @@
 from jax.ad_checkpoint import checkpoint_name
 import jax.numpy as jnp
 from jax.sharding import Mesh
-from MaxText import sharding
 from maxtext.common.common_types import Config, DecoderBlockType, EP_AS_CONTEXT, ShardMode
 from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN
 from maxtext.inference import page_manager
@@ -56,10 +55,11 @@
     simple_layer,
 )
 from maxtext.multimodal import utils as mm_utils
-from MaxText.sharding import create_sharding
+from maxtext.utils.sharding import create_sharding
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils
+from maxtext.utils import sharding
 
 # ------------------------------------------------------------------------------
 # The network: Decoder Definitions