feat(nnx): fix checkpoint restore ShapeDtypeStruct crash under native LoRA

.github/workflows/build_and_push_docker_image.yml

@@ -42,6 +42,10 @@ on:
         required: false
         type: boolean
         default: false
+      run_tests:
+        required: false
+        type: boolean
+        default: true
     outputs:
       build_result:
         description: 'The result of build_and_push job.'
@@ -184,7 +188,7 @@ jobs:
   run_ci_tests:
     name: Run Unit and Integration Tests
     needs: [pre_build_check, build_and_push]
-    if: needs.pre_build_check.outputs.should_run == 'true' && inputs.include_test_assets == true
+    if: needs.pre_build_check.outputs.should_run == 'true' && inputs.include_test_assets == true && inputs.run_tests == true
     uses: ./.github/workflows/run_ci_tests.yml
     with:
       image_name: ${{ inputs.image_name }}

.github/workflows/nightly_images_pipeline.yml

@@ -50,7 +50,7 @@ jobs:
       device_name: v4-8
       cloud_runner: linux-x86-n2-16-buildkit
 
-  build_and_push_docker_image:
+  build_and_push_tpu_docker_images:
     name: Build ${{ matrix.name }} Docker Image
     needs: build_maxtext_package
     strategy:
@@ -75,6 +75,26 @@ jobs:
             workflow: post-training
             image_name: maxtext_post_training_nightly
             dockerfile: maxtext_tpu_dependencies.Dockerfile
+    uses: ./.github/workflows/build_and_push_docker_image.yml
+    with:
+      image_name: ${{ inputs.image_suffix != '' && format('{0}_{1}', matrix.image_name, inputs.image_suffix) || matrix.image_name }}
+      device: ${{ matrix.device }}
+      build_mode: ${{ matrix.build_mode }}
+      workflow: ${{ matrix.workflow }}
+      dockerfile: ${{ matrix.dockerfile }}
+      maxtext_sha: ${{ needs.build_maxtext_package.outputs.maxtext_sha }}
+      include_test_assets: true
+      run_tests: false
+    secrets:
+      HF_TOKEN: ${{ secrets.HF_TOKEN }}
+
+  build_and_push_gpu_docker_images:
+    name: Build ${{ matrix.name }} Docker Image
+    needs: build_maxtext_package
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
           - name: "GPU Pre-Training Stable"
             device: gpu
             build_mode: stable
@@ -94,14 +114,40 @@ jobs:
       build_mode: ${{ matrix.build_mode }}
       workflow: ${{ matrix.workflow }}
       dockerfile: ${{ matrix.dockerfile }}
-      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
+      maxtext_sha: ${{ needs.build_maxtext_package.outputs.maxtext_sha }}
       include_test_assets: true
     secrets:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
 
+  run_tpu_ci_tests:
+    name: Run TPU ${{ matrix.name }} CI Tests
+    needs: build_and_push_tpu_docker_images
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - name: "Pre-Training Stable"
+            image_name: maxtext_jax_stable
+            device: tpu
+            workflow: pre-training
+          - name: "Pre-Training Nightly"
+            image_name: maxtext_jax_nightly
+            device: tpu
+            workflow: pre-training
+          - name: "Post-Training Nightly"
+            image_name: maxtext_post_training_nightly
+            device: tpu
+            workflow: post-training
+    uses: ./.github/workflows/run_ci_tests.yml
+    with:
+      image_name: ${{ inputs.image_suffix != '' && format('{0}_{1}', matrix.image_name, inputs.image_suffix) || matrix.image_name }}
+      image_tag: ${{ github.run_id }}
+      device: ${{ matrix.device }}
+      workflow: ${{ matrix.workflow }}
+
   run_e2e_tests:
     name: Run E2E tests
-    needs: build_and_push_docker_image
+    needs: build_and_push_tpu_docker_images
     uses: ./.github/workflows/run_e2e_tests.yml
     with:
       mode: ${{ inputs.image_suffix != '' && format('{0}_{1}', 'nightly', inputs.image_suffix) || 'nightly' }}
@@ -111,7 +157,7 @@ jobs:
 
   notify_failure:
     name: Notify failed build
-    needs: [build_and_push_docker_image, run_e2e_tests]
+    needs: [build_and_push_tpu_docker_images, build_and_push_gpu_docker_images, run_tpu_ci_tests, run_e2e_tests]
     if: ${{ failure() && inputs.image_suffix == '' }}
     runs-on: ubuntu-latest
     permissions:

.github/workflows/run_tests_coordinator.yml

@@ -209,7 +209,7 @@ jobs:
             "gpu-integration": "--ignore=tests/post_training",
             "cpu-unit": "--ignore=tests/post_training",
             "cpu-post-training-unit": "",
-            "cpu-torch-reference": "-o addopts= -rf --import-mode=importlib --strict-markers tests/unit/gemma4_layers_test.py tests/unit/gemma4_small_layers_test.py tests/unit/qwen3_next_vs_reference_test.py tests/unit/qwen3_5_layers_test.py"
+            "cpu-torch-reference": "-o addopts= -rf --import-mode=importlib --strict-markers tests/unit/gemma4_layers_test.py tests/unit/gemma4_small_layers_test.py tests/unit/qwen3_next_vs_reference_test.py tests/unit/qwen3_5_layers_test.py tests/unit/qwen3_omni_layers_test.py"
           }')[inputs.flavor] }}
         ${{ inputs.additional_pytest_args }}
 

scripts/diloco/run_diloco.sh

@@ -0,0 +1,101 @@
+#!/bin/bash
+
+# This script launches a DiLoCo pre-training workload on a GKE cluster using XPK.
+
+set -e
+
+# --- Cluster Parameters ---
+export PROJECT_ID="${PROJECT_ID:-}"
+export CLUSTER_NAME="${CLUSTER_NAME:-}"
+export ZONE="${ZONE:-}"
+export RESERVATION="${RESERVATION:-}"  # optional
+export BASE_OUTPUT_DIRECTORY="${BASE_OUTPUT_DIRECTORY:-}" # change to your own GCS bucket for logging and checkpointing
+export DATASET_PATH="${DATASET_PATH:-}" # change to your own GSC bucket for datasets. Make sure datasets exists
+export DOCKER_IMAGE="${DOCKER_IMAGE:-}" # Full path to the Docker image you pushed (e.g., gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:2026-06-22)
+export TPU_TYPE="${TPU_TYPE:-}"  # At least v5p-32 is needed to run Qwen3-30b-a3b. v5p-8 for qwen3-8b
+export NUM_SLICES="${NUM_SLICES:-}"  # you need at least two slices to let diloco take effect
+export WORKLOAD_NAME="${WORKLOAD_NAME:-$(whoami)-diloco-${TPU_TYPE}-$(date +%Y%m%d-%H%M%S)}" # this will be the name of run, for logging purposes
+
+# --- Model Parameters ---
+export MODEL_NAME="${MODEL_NAME:-}"
+export PER_DEVICE_BATCH_SIZE="${PER_DEVICE_BATCH_SIZE:-8}"
+export MAX_TARGET_LENGTH="${MAX_TARGET_LENGTH:-2048}"
+export TRAINING_STEPS="${TRAINING_STEPS:-20}"
+
+# --- DiLoCo Parameters ---
+export DILOCO_SYNC_PERIOD="${DILOCO_SYNC_PERIOD:-10}"
+export DILOCO_OUTER_LR="${DILOCO_OUTER_LR:-0.1}"
+export DILOCO_OUTER_MOMENTUM="${DILOCO_OUTER_MOMENTUM:-0.9}"
+
+# --- XLA Flags ---
+export XLA_FLAGS="${XLA_FLAGS:- \
+  --xla_tpu_scoped_vmem_limit_kib=65536 \
+--xla_tpu_bf16_emission_mode=NATIVE_EMISSION \
+--xla_tpu_enable_sparse_core_reduce_scatter_v2=true \
+--xla_tpu_enable_sparse_core_collective_offload_all_gather=true \
+--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true \
+--xla_tpu_enable_all_gather_offload_tracing=true \
+--xla_tpu_use_tc_device_shape_on_sc=True \
+--xla_sc_disable_megacore_partitioning=True \
+--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false \
+--xla_enable_async_all_gather=true \
+--xla_tpu_prefer_async_allgather_to_allreduce=true \
+--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \
+--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \
+--xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true \
+--xla_tpu_use_single_sparse_core_for_all_gather_offload=true \
+--xla_tpu_enable_concurrent_sparse_core_offloading=true \
+--xla_tpu_aggressive_opt_barrier_removal=true \
+--xla_tpu_enable_offloading_gather_to_sparsecore=true \
+--xla_tpu_sparse_core_all_gather_latency_multiplier=1 \
+--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 \
+--xla_tpu_enable_sparse_core_collective_aggregator=true \
+--xla_tpu_enable_latency_hiding_layer_scheduler=true \
+--xla_tpu_scheduler_percent_shared_memory_limit=150 \
+--xla_tpu_enable_layer_scheduler_for_dependent_collectives=true \
+--xla_tpu_enable_sparse_core_collective_offload_nd_reduce_scatter=true \
+--xla_tpu_pcie_bandwidth_multiplier=0.03 \
+--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true \
+--xla_tpu_enable_multi_compute_overlap_in_layer_scheduler=false \
+  --xla_tpu_enable_3d_reduce_scatter_decomposer=false }"
+
+if [ "$NUM_SLICES" -lt 2 ]; then
+    echo "Warning: NUM_SLICES is less than 2. DiLoCo will not take effect."
+fi
+
+# MaxText command
+MAXTEXT_COMMAND="cd /deps/src/ && \
+LIBTPU_INIT_ARGS='${XLA_FLAGS}' \
+python3 -m maxtext.trainers.pre_train.train \
+maxtext/configs/base.yml \
+run_name=$WORKLOAD_NAME \
+save_config_to_gcs=true \
+base_output_directory=$BASE_OUTPUT_DIRECTORY \
+dataset_path=$DATASET_PATH \
+dataset_name='c4/en:3.0.1' \
+eval_dataset_name='c4/en:3.0.1' \
+model_name=$MODEL_NAME \
+tokenizer_type=huggingface \
+tokenizer_path=maxtext/assets/tokenizers/qwen3-tokenizer \
+per_device_batch_size=$PER_DEVICE_BATCH_SIZE \
+max_target_length=$MAX_TARGET_LENGTH \
+enable_diloco=true \
+dcn_diloco_parallelism=$NUM_SLICES \
+diloco_sync_period=$DILOCO_SYNC_PERIOD \
+diloco_outer_lr=$DILOCO_OUTER_LR \
+diloco_outer_momentum=$DILOCO_OUTER_MOMENTUM \
+steps=$TRAINING_STEPS"
+
+# Workload Creation
+xpk workload create \
+  --cluster="$CLUSTER_NAME" \
+  --project="$PROJECT_ID" \
+  --zone="$ZONE" \
+  --priority=medium \
+  --max-restarts=0 \
+  --tpu-type="$TPU_TYPE" \
+  --num-slices="$NUM_SLICES" \
+  --docker-image="${DOCKER_IMAGE}" \
+  --workload="${WORKLOAD_NAME}" \
+  ${RESERVATION:+--reservation="$RESERVATION"} \
+  --command="${MAXTEXT_COMMAND}"

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

@@ -117,7 +117,7 @@ def init_state_fn():
       cfg.checkpoint_period,
   )
 
-  state, _, _, _ = maxtext_utils.setup_training_state(None, cfg, mesh, checkpoint_manager, init_state_fn)
+  state, _, _, _, _ = maxtext_utils.setup_training_state(None, cfg, mesh, checkpoint_manager, init_state_fn)
   max_logging.log("start")
   max_utils.print_mem_stats("After params initialized")
 

src/maxtext/common/checkpointing.py

@@ -202,19 +202,27 @@ def _populate_pure_dict_from_partial(abstract_pure, partial_concrete, config=Non
   (NNX-only state the Linen checkpoint never had) get `_default_for_sds` or are preserved.
   """
   if isinstance(abstract_pure, dict):
-    return {
-        k: _populate_pure_dict_from_partial(
-            v,
-            partial_concrete.get(k) if isinstance(partial_concrete, dict) else None,
-            config,
-        )
-        for k, v in abstract_pure.items()
-    }
+    res = {}
+    for k, v in abstract_pure.items():
+      val = _populate_pure_dict_from_partial(
+          v,
+          partial_concrete.get(k) if isinstance(partial_concrete, dict) else None,
+          config,
+      )
+      if val is not None:
+        res[k] = val
+    return res
+
+  if isinstance(partial_concrete, jax.ShapeDtypeStruct):
+    partial_concrete = None
+
   if partial_concrete is not None and not isinstance(partial_concrete, dict):
     return partial_concrete
 
-  # Under LoRA, preserve the already-initialized base weights instead of resetting to zeros
+  # Under LoRA, return None for SDS to omit from state dict and preserve base weights.
   if config and config.lora and config.lora.enable_lora:
+    if isinstance(abstract_pure, jax.ShapeDtypeStruct):
+      return None
     return abstract_pure
 
   return _default_for_sds(abstract_pure)

src/maxtext/configs/base.yml

@@ -881,7 +881,7 @@ diloco_sync_period: 36
 diloco_outer_lr: 0.3
 diloco_outer_momentum: 0.9
 # DCN bandwidth throttling parameters (used for simulating slow networks).
-# If dcn_bandwidth_limit is empty, no throttling is applied.
+# If dcn_bandwidth_limit is empty, no throttling is applied. **Per-VM** throughput limit. Example: 10gbit.
 dcn_bandwidth_limit: ""
 # The burst size parameter for the traffic control (tc) token bucket filter.
 dcn_bandwidth_burst: "10mb"

src/maxtext/experimental/rl/grpo_trainer.py

@@ -864,7 +864,7 @@ def init_state_fn():
 
   with maybe_record_goodput(recorder, GoodputEvent.TRAINING_PREPARATION):
     data_iterator = grpo_input_pipeline.create_data_iterator(config_inference, inference_mesh)
-    state, _, state_mesh_shardings, data_iterator = maxtext_utils.setup_training_state(
+    state, _, state_mesh_shardings, data_iterator, _ = maxtext_utils.setup_training_state(
         data_iterator, config, mesh, checkpoint_manager, init_state_fn
     )