FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/AI-Hypercomputer/maxtext/pull/4513 from Shuwen-Fang:sort_one_sparsecore cccf4ecd7360fbb63557537ccc731eb4f5d0b102

.github/workflows/gemini_invoke.yml

@@ -48,7 +48,7 @@
         id: 'mint_identity_token'
         if: |-
           ${{ vars.APP_ID }}
         uses: 'actions/create-github-app-token@v2'
         with:
           app-id: '${{ vars.APP_ID }}'
           private-key: '${{ secrets.APP_PRIVATE_KEY }}'
@@ -59,7 +59,7 @@
       - name: 'Run Gemini CLI'
       # Trigger Gemini with context
         id: 'run_gemini'
         uses: 'google-github-actions/run-gemini-cli@main'
         env:
           GEMINI_CLI_TRUST_WORKSPACE: 'true'
           TITLE: '${{ github.event.pull_request.title || github.event.issue.title }}'
@@ -86,7 +86,7 @@
           settings: |-
             {
               "model": {
-                "maxSessionTurns": 25
+                "maxSessionTurns": 100
               },
               "telemetry": {
                 "enabled": false,

.github/workflows/gemini_review.yml

@@ -36,7 +36,7 @@
 jobs:
   review:
     runs-on: 'ubuntu-latest'
-    timeout-minutes: 10
+    timeout-minutes: 30
     permissions:
       contents: 'read'
       id-token: 'write'
@@ -47,7 +47,7 @@
         id: 'mint_identity_token'
         if: |-
           ${{ vars.APP_ID }}
         uses: 'actions/create-github-app-token@v2'
         with:
           app-id: '${{ vars.APP_ID }}'
           private-key: '${{ secrets.APP_PRIVATE_KEY }}'
@@ -57,7 +57,7 @@
 
       - name: 'Checkout repository'
         # downloads the code to be analyzed
         uses: 'actions/checkout@v6'
         with:
           persist-credentials: 'false'
 
@@ -77,7 +77,7 @@
 
       - name: 'Run Gemini pull request review'
         # reviews code with detailed set of instructions for the Gemini
         uses: 'google-github-actions/run-gemini-cli@v0'
         id: 'gemini_pr_review'
         env:
           GEMINI_CLI_TRUST_WORKSPACE: 'true'
@@ -103,7 +103,7 @@
           settings: |-
             {
               "model": {
-                "maxSessionTurns": 25
+                "maxSessionTurns": 100
               },
               "telemetry": {
                 "enabled": false,

.github/workflows/promote_docker_image.yml

@@ -40,12 +40,14 @@ jobs:
           DAG_RUN_ID: ${{ github.event.client_payload.dag_run_id }}
           SHA: ${{ github.event.client_payload.sha }}
           GITHUB_RUN_ID: ${{ github.event.client_payload.github_run_id }}
+          TEST_TYPE: ${{ github.event.client_payload.test_type }}
         run: |
           echo "================================"
           echo "Github Run ID: ${GITHUB_RUN_ID}"
           echo "DAG ID:        ${DAG_ID}"
           echo "DAG Run ID:    ${DAG_RUN_ID}"
           echo "Commit SHA:    ${SHA}"
+          echo "Test Type:     ${TEST_TYPE}"
           echo "State:         ${STATE}"
           echo "================================"
 
@@ -66,25 +68,28 @@ jobs:
   tag_docker_image:
     name: Promote ${{ matrix.image_name }} Docker Image
     needs: handle_result
-    if: needs.handle_result.result == 'success'
+    if: needs.handle_result.result == 'success' && (github.event.client_payload.test_type == '' || github.event.client_payload.test_type == matrix.test_type)
     runs-on: linux-x86-n2-16-buildkit
     container: google/cloud-sdk:524.0.0
     strategy:
       fail-fast: false
       matrix:
-        image_name:
-          - maxtext_jax_nightly
-          - maxtext_gpu_jax_nightly
-          - maxtext_post_training_nightly
+        include:
+          - test_type: pre_training
+            image_name: maxtext_jax_nightly
+          - test_type: post_training
+            image_name: maxtext_post_training_nightly
     steps:
       - name: Configure Docker
         run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q
       - name: Add tags to Docker image
         shell: bash
         env:
           GITHUB_RUN_ID: ${{ github.event.client_payload.github_run_id }}
+          PROJECT_NAME: ${{ vars.PROJECT_NAME }}
+          IMAGE_NAME: ${{ matrix.image_name }}
         run: |
-          SOURCE_IMAGE="gcr.io/${{ vars.PROJECT_NAME }}/${{ matrix.image_name }}"
+          SOURCE_IMAGE="gcr.io/${PROJECT_NAME}/${IMAGE_NAME}"
 
           # Add the traceability tag to confirm it passed validation suite
           gcloud container images add-tag "${SOURCE_IMAGE}:${GITHUB_RUN_ID}" \

src/maxtext/checkpoint_conversion/inspect_checkpoint.py

@@ -48,7 +48,7 @@
 
   [Mode 2: MaxText Architecture]
     python -m maxtext.checkpoint_conversion.inspect_checkpoint maxtext \
-        model_name=<maxtext_model_name> scan_layers=<True | False> 
+        model_name=<maxtext_model_name> scan_layers=<True | False> enable_nnx=<True | False>
         (Optional: other maxtext config)
 
   [Mode 3: Orbax]
@@ -68,6 +68,8 @@
 import absl
 from maxtext.inference.inference_utils import str2bool
 from maxtext.checkpoint_conversion.utils.utils import print_peak_memory
+from maxtext.utils.model_creation_utils import create_nnx_abstract_model
+from flax import nnx
 
 
 def natural_sort_key(s: str):
@@ -238,15 +240,19 @@ def inspect_maxtext(args, remaining_args):
   print(argv)
   config = pyconfig.initialize(argv)
 
-  print(f"\n--- Inspecting MaxText Architecture: {config.model_name} (Scan: {config.scan_layers}) ---")
+  print(
+      f"\n--- Inspecting MaxText Architecture: {config.model_name} "
+      f"(scan_layers: {config.scan_layers}, enable_nnx: {config.enable_nnx}) ---"
+  )
   devices_array = maxtext_utils.create_device_mesh(config)
   mesh = jax.sharding.Mesh(devices_array, config.mesh_axes)
-  quant = quantizations.configure_quantization(config)
-  model = Transformer(config, mesh=mesh, quant=quant)
-
-  # Extract abstract parameters. This returns a PyTree of `ShapeDtypeStruct`
-  # objects, without materializing weights.
-  abstract_param = maxtext_utils.get_abstract_param(model, config)
+  if config.enable_nnx:
+    _, abstract_model = create_nnx_abstract_model(config, mesh=mesh)
+    _, abstract_param, _ = nnx.split(abstract_model, nnx.Param, ...)
+  else:
+    quant = quantizations.configure_quantization(config)
+    model = Transformer(config, mesh=mesh, quant=quant)
+    abstract_param = maxtext_utils.get_abstract_param(model, config)
 
   # Calculate and display the total parameter count based purely on abstract shapes.
   num_params = max_utils.calculate_num_params_from_pytree(abstract_param)
@@ -261,8 +267,14 @@ def inspect_maxtext(args, remaining_args):
   for path_tuple, abstract_leaf_value in abstract_params_flat:
     key_parts = param_key_parts_from_path(path_tuple)
 
-    # Construct a MaxText-style parameter key (e.g., "params.params.layer.weight").
-    param_key = "params." + ".".join(key_parts)
+    # Construct a MaxText-style parameter key. Examples:
+    # "params.params.decoder.decoder_norm.scale" (for standard model weights)
+    # "params.Tid2EidVar.decoder.layers_0.mlp.MoeBlock_0.tid2eid" (for legacy custom collections)
+    key_str = ".".join(key_parts)
+    if config.enable_nnx and not key_str.startswith(("params", "Tid2EidVar")):
+      param_key = "params.params." + key_str
+    else:
+      param_key = "params." + key_str
 
     shape = abstract_leaf_value.shape
     param_dict[param_key] = f"shape: {shape}"
@@ -294,6 +306,7 @@ def inspect_orbax(args):
   # Defer imports to avoid overhead when running in other modes.
   import orbax.checkpoint as ocp
   from etils import epath
+  from maxtext.checkpoint_conversion.utils.utils import param_key_parts_from_path
 
   path = epath.Path(args.path)
 
@@ -309,9 +322,13 @@ def inspect_orbax(args):
   # Filter strictly for parameter keys and format them.
   param_dict = {}
   for k, v in dictionary.items():
-    # `k` is a tuple representing the path hierarchy; join it into a single string.
-    # `v` is a metadata object containing `.shape` and `.dtype`.
-    param_key = ".".join(k)
+    # `k` is a tuple representing the path hierarchy.
+    # `v` is a shape-dtype struct.
+    # Pass it through param_key_parts_from_path to normalize NNX structures
+    # (folding list indices like "0" into "layers_0" and dropping "value").
+    key_parts = param_key_parts_from_path(k)
+    param_key = ".".join(key_parts)
+
     if not param_key.startswith("params"):
       continue
 
@@ -372,7 +389,8 @@ def main():
   if args.mode == "hf":
     inspect_hf(args)
   elif args.mode == "maxtext":
-    # remaining_args accepts maxtext config, like `model_name=<maxtext_model_name> scan_layers=<True | False>`
+    # remaining_args accepts maxtext config, like `model_name=<maxtext_model_name>
+    # scan_layers=<True | False> enable_nnx=<True | False>`
     inspect_maxtext(args, remaining_args)
   elif args.mode == "orbax":
     inspect_orbax(args)

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

@@ -41,6 +41,7 @@
 import sys
 
 from flax import nnx
+from flax import serialization
 import jax
 from jax import random
 from jax.sharding import Mesh
@@ -118,6 +119,12 @@ def init_state_fn():
   )
 
   state, _, _, _, _ = maxtext_utils.setup_training_state(None, cfg, mesh, checkpoint_manager, init_state_fn)
+  if cfg.pure_nnx:
+    state = train_state_nnx.to_checkpoint_dict(state)
+    state.pop("nnx_aux", None)
+  else:
+    state = serialization.to_state_dict(state)
+
   max_logging.log("start")
   max_utils.print_mem_stats("After params initialized")
 
@@ -201,24 +208,10 @@ def init_state_fn():
       "['decoder']['decoder_norm']['bias']": (".params.lm.final_ln.bias", None),
   }
 
-  if cfg.pure_nnx:
-    # NNX state-tree paths after `nnx.split(TrainStateNNX)`. The state is a
-    # nested `nnx.State` (dict-like Mapping) with `nnx.Variable` leaves, so
-    # `jax.tree_util.keystr` produces dict-style entries (`['key']`) plus
-    # `.value` for the Variable leaf, plus `[idx]` for the optax tuple:
-    #   model params   -> ['model']<rest>.value
-    #   adam mu / nu   -> ['optimizer']['opt_state'][0]['mu' | 'nu']<rest>.value
-    #   step           -> ['optimizer']['step'].value
-    #   opt count      -> ['optimizer']['opt_state'][0]['count'].value
-    state_map = {
-        "['optimizer']['step'].value": ("step", None),
-        "['optimizer']['opt_state'][0]['count'].value": ("opt_states_0.no_prefix_0.count", None),
-    }
-  else:
-    state_map = {
-        ".step": ("step", None),
-        ".opt_state.count": ("opt_states_0.no_prefix_0.count", None),
-    }
+  state_map = {
+      "['step']": ("step", None),
+      "['opt_state']['count']": ("opt_states_0.no_prefix_0.count", None),
+  }
 
   def get_layer_prefix(keystr_pax):
     # different path format between decoder_layer variable
@@ -231,26 +224,18 @@ def get_layer_prefix(keystr_pax):
 
   for keystr_maxtext, (keystr_pax, transform_fn) in keystr_map.items():
     prefix_pax_opt_state = get_layer_prefix(keystr_pax)
-    if cfg.pure_nnx:
-      state_map[f"['model']{keystr_maxtext}.value"] = (f"mdl_vars{keystr_pax}", transform_fn)
-      state_map[f"['optimizer']['opt_state'][0]['mu']{keystr_maxtext}.value"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
-          transform_fn,
-      )
-      state_map[f"['optimizer']['opt_state'][0]['nu']{keystr_maxtext}.value"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
-          transform_fn,
-      )
-    else:
-      state_map[f".params['params']{keystr_maxtext}"] = (f"mdl_vars{keystr_pax}", transform_fn)
-      state_map[f".opt_state.mu['params']{keystr_maxtext}"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
-          transform_fn,
-      )
-      state_map[f".opt_state.nu['params']{keystr_maxtext}"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
-          transform_fn,
-      )
+    state_map[f"['params']['params']{keystr_maxtext}"] = (
+        f"mdl_vars{keystr_pax}",
+        transform_fn,
+    )
+    state_map[f"['opt_state']['mu']['params']{keystr_maxtext}"] = (
+        f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
+        transform_fn,
+    )
+    state_map[f"['opt_state']['nu']['params']{keystr_maxtext}"] = (
+        f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
+        transform_fn,
+    )
 
   def verify_fn(key_path, _):
     keystr = jax.tree_util.keystr(key_path)
@@ -302,9 +287,9 @@ def map_fn(key_path, value):
   max_logging.log("converted state finished")
   max_utils.print_mem_stats("converted state finished")
 
-  step_value = int(converted_state.optimizer.step.value) if cfg.pure_nnx else converted_state.step
-  if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state):
-    max_logging.log(f"saved a checkpoint at step {step_value}")
+  step_value = int(converted_state["step"]) if isinstance(converted_state, dict) else int(converted_state.step)
+  if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state, config=cfg):
+    max_logging.log(f"Saved a checkpoint at step {step_value}.")
   # Upon preemption, exit when and only when all ongoing saves are complete.
   if checkpointing.reached_preemption(checkpoint_manager, step_value):
     checkpointing.wait_until_finished(checkpoint_manager)
@@ -322,12 +307,22 @@ def map_fn(key_path, value):
       default="gs://mlperf-llm-public2/gpt3_spmd1x64x24_tpuv4-3072_v84_20221101/checkpoints/checkpoint_00004000",
       required=True,
   )
-  parser.add_argument("--maxtext-model-name", choices=["gpt3-175b", "gpt3-52k"], type=str, required=True)
+  parser.add_argument(
+      "--maxtext-model-name",
+      choices=["gpt3-175b", "gpt3-52k"],
+      type=str,
+      required=True,
+  )
   parser.add_argument("--base-output-directory", type=str, required=True)
   parser.add_argument("--run-name", type=str, required=True)
 
   args = parser.parse_args()
   if not args.paxml_ckpt_path.startswith("gs://"):
     raise ValueError("--paxml-ckpt-path should be a gcs path starting with gs://")
 
-  convert(args.paxml_ckpt_path, args.maxtext_model_name, args.base_output_directory, args.run_name)
+  convert(
+      args.paxml_ckpt_path,
+      args.maxtext_model_name,
+      args.base_output_directory,
+      args.run_name,
+  )

src/maxtext/checkpoint_conversion/utils/load_dynamic.py

@@ -73,19 +73,23 @@
 
 from flax import nnx
 import flax.traverse_util
-from google.cloud import storage
 import huggingface_hub
 import jax
 from maxtext.checkpoint_conversion.utils import hf_model_configs
 from maxtext.checkpoint_conversion.utils import param_mapping
 from maxtext.checkpoint_conversion.utils import tensor_handling
+from maxtext.common.gcloud_stub import gcs_storage
 from maxtext.utils import gcs_utils
 from maxtext.utils import globals as maxtext_globals
 from maxtext.utils import max_logging
 from orbax.checkpoint import v1 as ocp_v1
 from orbax.checkpoint._src.arrays import sharding as sharding_utils
 
 
+# Route GCS through the decoupling helper so this module imports cleanly in
+# decoupled environments where google-cloud-storage is intentionally absent.
+storage = gcs_storage()
+
 HF_MODEL_CONFIGS = hf_model_configs.HF_MODEL_CONFIGS
 get_hf_loading_function = tensor_handling.get_hf_loading_function
 

src/maxtext/checkpoint_conversion/utils/utils.py

@@ -867,12 +867,26 @@ def create_restore_args(tree_metadata):
   for i, path in enumerate(paths):
     checkpoint_path = epath.Path(path)
     metadata = ckptr.metadata(checkpoint_path)
+    checkpoint_tree = metadata.item_metadata.tree
+    if isinstance(checkpoint_tree, dict):
+      if "params" in checkpoint_tree:
+        checkpoint_tree = {"params": checkpoint_tree["params"]}
+        max_logging.log(f"Filtering checkpoint to only load 'params' from {path}")
+      else:
+        filtered_tree = {k: v for k, v in checkpoint_tree.items() if k not in ("opt_state", "optimizer")}
+        if len(filtered_tree) < len(checkpoint_tree):
+          checkpoint_tree = filtered_tree
+          max_logging.log(f"Filtering checkpoint to exclude optimizer keys from {path}")
+
     restore_args = jax.tree_util.tree_map(
         lambda x: create_restore_args(x) if hasattr(x, "shape") else None,
-        metadata.item_metadata.tree,
+        checkpoint_tree,
         is_leaf=lambda x: hasattr(x, "shape"),
     )
-    restored = ckptr.restore(checkpoint_path, restore_args=restore_args)
+    restored = ckptr.restore(
+        checkpoint_path,
+        args=ocp.args.PyTreeRestore(item=checkpoint_tree, restore_args=restore_args, partial_restore=True),
+    )
 
     if i == 0:
       merged_dict = restored