Add a new goodput elastic utils function record_slice_state()

Author: YixuanWang-99

.github/workflows/run_pathways_tests.yml

@@ -72,11 +72,9 @@ jobs:
       options: ${{ inputs.container_resource_option }}
     steps:
       - name: Checkout MaxText
-        run: |
-          git config --global --add safe.directory /__w/maxtext/maxtext
-          git clone https://github.com/google/maxtext.git .
-          git fetch origin ${{ inputs.maxtext_sha }}
-          git checkout FETCH_HEAD
+        uses: actions/checkout@v5
+        with:
+          ref: ${{ inputs.maxtext_sha }}
       - name: Download the maxtext wheel
         uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5.0.0
         with:

README.md

@@ -42,7 +42,6 @@ See our guide on running MaxText in decoupled mode, without any GCP dependencies
 
 ## 🔥 Latest news 🔥
 
-- [June 9, 2026] Qwen3.5 35B & 397B is now [supported](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/qwen/moe/run_qwen_moe.md).
 - [May 29, 2026] A new eval framework is now available for running lm-eval, evalchemy, and custom benchmarks against MaxText checkpoints. See the [Evaluation guide](https://maxtext.readthedocs.io/en/latest/guides/eval_framework.html) to get started.
 - [May 6, 2026] Kimi-K2-Thinking, Kimi-K2.5 (text), and Kimi-K2.6 (text) are now supported. See [Run_Kimi.md](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/kimi/Run_Kimi.md#quantized-variants-k2-thinking-k25-k26) for details.
 - [April 18, 2026] Added a new notebook [maxtext_with_gepa.ipynb](https://github.com/AI-Hypercomputer/maxtext/blob/3c7d8d27864fc12cccac07786f02bd0e5262c982/src/maxtext/examples/maxtext_with_gepa.ipynb) for optimizing AIME prompts using the GEPA framework with Maxtext.
@@ -123,7 +122,6 @@ MaxText aims to provide you with the best OSS models, whether as a reference imp
   - Gemma 2 (2B, 9B, 27B)
   - Gemma 1 (2B, 7B)
 - Alibaba
-  - Qwen 3.5 MoE (35B, 397B)
   - Qwen 3 Next (80B)
   - Qwen 3 MoE (30B, 235B), Qwen 3 MoE 2507 (235B, 480B)
   - Qwen 3 Dense (0.6B, 1.7B, 4B, 8B, 14B, 32B)

docs/guides/checkpointing_solutions/convert_checkpoint.md

@@ -36,6 +36,10 @@ Use the `to_maxtext.py` script to convert a Hugging Face model checkpoint into a
 ### Setup Environment
 
 ```bash
+# Install PyTorch and safetensors (in MaxText virtual environment)
+python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+python3 -m pip install safetensors --no-deps
+
 # Setup environment variables
 export MODEL=<HF_MODEL> # e.g. 'llama3.1-8b-Instruct'
 export BASE_OUTPUT_DIRECTORY=<CKPT_PATH> # e.g., gs://my-bucket/my-checkpoint-directory
@@ -83,6 +87,10 @@ Use the `to_huggingface.py` script to convert a MaxText checkpoint into the Hugg
 ### Setup Environment
 
 ```bash
+# Install PyTorch and safetensors (in MaxText virtual environment)
+python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
+python3 -m pip install safetensors --no-deps
+
 # Setup environment variables
 export MODEL=<MODEL_NAME> # e.g. 'qwen3-4b'
 export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/0/items
@@ -128,9 +136,6 @@ To ensure the conversion was successful, you can use the [test script](https://g
 export MODEL=<MODEL_NAME> # e.g. 'qwen3-4b'
 export MAXTEXT_CKPT_PATH=<CKPT_PATH> # e.g., gs://my-bucket/my-model-checkpoint/0/items
 export HF_CKPT_PATH=<HF_CKPT_PATH> # e.g., gs://my-bucket/my-checkpoint-directory
-
-# Install PyTorch (in MaxText virtual environment)
-python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu
 ```
 
 ### Run Correctness Test

docs/reference/models/supported_models_and_architectures.md

@@ -47,11 +47,6 @@ MaxText is an open-source, high-performance LLM framework written in Python/JAX.
 - **Variants**: Dense (0.6B–32B); MoE (30B-A3B, 235B-A22B, 480B Coder), MoE with Hybrid Attention (Next-80B-A3B)
 - **Notes**: **QK-Norm**, GQA, SwiGLU, RMSNorm, RoPE, GatedDeltaNet.
 
-### Qwen3.5
-
-- **Variants** MoE with Hybrid Attention (35B-A3B, 397B-A17B)
-- **Notes**: GatedDeltaNet, GQA, Multimodal, MoE.
-
 ### GPT-OSS
 
 - **Variants**: 20B, 120B
@@ -94,10 +89,7 @@ The following summarizes observed runtime efficiency and scaling behaviors of Ma
   - **Gemma**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/gemma/Run_Gemma.md) | [Gemma Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/gemma.py) | [Gemma2 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/gemma2.py) | [Gemma3 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/gemma3.py) | [Gemma4 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/gemma4.py)
   - **Mixtral**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/mixtral/Run_Mixtral.md) | [Mixtral Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/mixtral.py) | [Mistral Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/mistral.py)
   - **DeepSeek**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/deepseek/Run_DeepSeek.md) | [DeepSeek Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/deepseek.py)
-  - **Qwen**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/qwen/moe/run_qwen_moe.md) |
-    [Qwen3-Next Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/qwen/next/run_qwen3_next.md) |
-    [Qwen3 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/qwen3.py) | [Qwen3-Next Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/qwen3.py) |
-    [Qwen3.5 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/qwen3_5.py)
+  - **Qwen3**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/qwen/moe/run_qwen_moe.md) | [Qwen3-Next Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/qwen/next/run_qwen3_next.md) | [Qwen3 Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/qwen3.py) | [Qwen3-Next Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/qwen3.py)
   - **GPT-OSS**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/gpt_oss/run_gpt_oss.md) | [GPT-OSS Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/gpt_oss.py)
   - **Kimi**: [Guide](https://github.com/AI-Hypercomputer/maxtext/blob/main/tests/end_to_end/tpu/kimi/Run_Kimi.md) | [K2 reuses DeepSeek Source](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/deepseek.py)
 

docs/release_notes.md

@@ -22,37 +22,6 @@ MaxText is [available in PyPI](https://pypi.org/project/maxtext/) and can be ins
 
 ## Releases
 
-### v0.2.3
-
-#### Changes
-
-- Upgraded JAX to version 0.10.0 for pre-training and 0.10.1 for post-training.
-- **New vLLM-Powered Evaluation Framework**: Introduced an eval framework for running lm-eval, evalchemy, and custom benchmarking against MaxText checkpoints. See the [evaluation guide](https://maxtext.readthedocs.io/en/latest/guides/eval_framework.html) for details.
-- Added support for pre-training new models:
-  - **Qwen3.5**: Qwen3.5 35B & 397B is now [supported](https://github.com/AI-Hypercomputer/maxtext/blob/d938b91acaa3baaaf32956e21677bd29e14549a1/tests/end_to_end/tpu/qwen/moe/run_qwen_moe.md).
-  - **Qwen3-Omni**: Support for multimodal SFT ([PR #3863](https://github.com/AI-Hypercomputer/maxtext/pull/3863)).
-- **Direct Preference Optimization (DPO/ORPO) Support**: Full support for DPO and ORPO alignment pipelines. See the [DPO tutorial](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/dpo.html) for details.
-- **Reinforcement Learning (RL) Recipe**: Added a pre-configured [RL recipe for Qwen3-30b-a3b](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/rl_qwen3_30b.html).
-- **Iterative Quality Monitoring (RL)**: Added intermediate evaluation hooks to automatically run quality benchmarks during RL training (every `eval_interval` steps), optimized with a new `eval_batch_size` configuration knob.
-- **Developer Extensibility**: Added `dataset_processor_path` CLI knob for custom dataset integration, and refactored shared post-training hooks to simplify custom SFT, DPO, and RL workflow development.
-- **Generalized Learn-to-Init (LTI) for Distillation**: Enhanced post-training distillation capabilities with generalized LTI support.
-- Added support for recording elastic goodput events during training to track efficiency ([PR #3901](https://github.com/AI-Hypercomputer/maxtext/pull/3901)).
-- **Installation Updates**: Updated the `[tpu-post-train]` installation command to require `UV_TORCH_BACKEND=cpu`(see [Installation Guide](install_maxtext.md)).
-- **Zero1 AOT Compilation**: Added zero1 support to Ahead-Of-Time (AOT) compilation in train compile, improving compilation capabilities for zero1 config.
-- **MoE Performance Optimization**: Integrated ragged gather reduce into Mixture of Experts (MoE) layers to optimize memory and performance by replacing ragged scatter and supporting backward pass.
-- Added [E2E scripts](https://github.com/AI-Hypercomputer/maxtext/tree/main/tests/end_to_end/tpu/gemma3/4b) to run checkpoint conversion, pre-training and post-training (SFT, RL) with Gemma3-4B model.
-- **Bug Fixes and Usability Enhancements**:
-  - **Attention Masking Fix in RL**: Fixed an issue in `TunixMaxTextAdapter` where queries at non-pad positions could attend to pad-position keys during training, which was corrupting log-probabilities and affecting GRPO training reward trajectories ([PR #4016](https://github.com/AI-Hypercomputer/maxtext/pull/4016)).
-  - **JAX/NNX Gradient Mutation Fix**: Refactored post-training loops (`train_distill`, `train_sft`, `train_rl`) to use `jax.value_and_grad` with explicit NNX state split/merge instead of nesting `nnx.value_and_grad` inside `nnx.jit` ([PR #3652](https://github.com/AI-Hypercomputer/maxtext/pull/3652)).
-  - **Qwen3-MoE Checkpoint Conversion**: Fixed checkpoint conversion issues for Qwen3-MoE models ([PR #3868](https://github.com/AI-Hypercomputer/maxtext/pull/3868)).
-  - **Duplicate Configuration Failures Fix**: Allowed identical config overrides and handled configuration exceptions cleanly ([PR #3933](https://github.com/AI-Hypercomputer/maxtext/pull/3933)).
-- **Documentation Improvements**: Updated [Getting started](https://maxtext.readthedocs.io/en/latest/getting_started.html) guide, including new guides for the [evaluation framework](https://maxtext.readthedocs.io/en/latest/guides/eval_framework.html) and the [DPO tutorial](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/dpo.html).
-
-#### Deprecations
-
-- Deleted [legacy DPO implementation](https://github.com/AI-Hypercomputer/maxtext/pull/3997) in favor of the integrated [DPO trainer](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/dpo.html).
-- Removed stack trace collection feature.
-
 ### v0.2.2
 
 #### Changes

src/maxtext/__init__.py

@@ -19,7 +19,7 @@
 """
 
 __author__ = "Google LLC"
-__version__ = "0.2.3"
+__version__ = "0.2.2"
 __description__ = (
     "MaxText is a high performance, highly scalable, open-source LLM written in pure Python/Jax and "
     "targeting Google Cloud TPUs and GPUs for training and **inference."

src/maxtext/checkpoint_conversion/to_maxtext.py

@@ -67,7 +67,7 @@
 from maxtext.common.common_types import MODEL_MODE_TRAIN
 from maxtext.checkpoint_conversion.utils.hf_model_configs import HF_MODEL_CONFIGS
 from maxtext.checkpoint_conversion.utils.param_mapping import HOOK_FNS, PARAM_MAPPING
-from maxtext.checkpoint_conversion.utils.utils import MemoryMonitorTqdm, apply_hook_fns, load_hf_dict_from_transformers, load_hf_dict_from_safetensors, param_key_parts_from_path, print_peak_memory, print_ram_usage, save_weights_to_checkpoint, validate_and_filter_param_map_keys
+from maxtext.checkpoint_conversion.utils.utils import MemoryMonitorTqdm, apply_hook_fns, load_hf_dict_from_transformers, load_hf_dict_from_safetensors, print_peak_memory, print_ram_usage, save_weights_to_checkpoint, validate_and_filter_param_map_keys
 from maxtext.inference.inference_utils import str2bool
 from maxtext.layers import quantizations
 from maxtext.models import models
@@ -333,7 +333,8 @@ def get_maxtext_model_info(config):
   # preprocess state
   maxtext_abstract_dict = {}
   for mt_target_idx, (path_tuple, abstract_leaf_value) in enumerate(abstract_params_flat):
-    mt_param_key = "params-" + "-".join(param_key_parts_from_path(path_tuple))
+    key_parts = [k.key for k in path_tuple if hasattr(k, "key")]
+    mt_param_key = "params-" + "-".join(key_parts)
     mt_target_shape = abstract_leaf_value.shape
     maxtext_abstract_dict[mt_param_key] = (mt_target_idx, mt_target_shape)
 

src/maxtext/checkpoint_conversion/utils/hf_shape.py

@@ -1131,7 +1131,6 @@ def MIXTRAL_HF_WEIGHTS_TO_SHAPE(config):
     "llama3.1-70b": LLAMA31_HF_WEIGHTS_TO_SHAPE,
     "llama3.1-405b": LLAMA31_HF_WEIGHTS_TO_SHAPE,
     "qwen3-30b-a3b": QWEN_HF_WEIGHTS_TO_SHAPE,
-    "qwen3-30b-a3b-base": QWEN_HF_WEIGHTS_TO_SHAPE,
     "qwen3-235b-a22b": QWEN_HF_WEIGHTS_TO_SHAPE,
     "qwen3-480b-a35b": QWEN_HF_WEIGHTS_TO_SHAPE,
     "deepseek2-16b": DEEPSEEK_HF_WEIGHTS_TO_SHAPE,

src/maxtext/checkpoint_conversion/utils/utils.py

@@ -888,51 +888,6 @@ def save_adapter_files(output_dir, weights, config, found_modules, model_id):
     json.dump(adapter_config, f, indent=4)
 
 
-def param_key_parts_from_path(path_tuple) -> list[str]:
-  """Convert a JAX tree path into MaxText dash-joined key segments.
-
-  Normalizes two NNX storage artifacts so the result follows the MaxText-Linen
-  naming convention and matches the param-mapping tables (e.g.
-  ``params-decoder-layers_0-self_attention-query-kernel``):
-
-  * ``nnx.List`` layer stacks flatten to an *integer* path key
-    (``decoder -> layers -> 0 -> ...``), which Orbax may restore as a numeric
-    *string* (``"0"``). Either form is folded into the preceding segment as
-    ``<name>_<idx>`` (``layers_0``), matching Linen's ``layers_0`` name. (A pure
-    integer key would otherwise raise ``TypeError: sequence item N: expected str
-    instance, int found`` when joined; a string ``"0"`` would mismatch the
-    ``layers_0`` mapping.)
-  * ``nnx.Variable`` leaves flatten with a trailing ``value`` key
-    (``...-kernel -> value``). That wrapper segment is dropped, since MaxText-Linen
-    param keys have no such suffix.
-
-  Scanned / plain Linen string paths (no integer key, no trailing ``value``) are
-  returned unchanged.
-
-  Args:
-    path_tuple: A path produced by ``jax.tree_util.tree_flatten_with_path`` or
-      ``tree_leaves_with_path`` (a sequence of ``DictKey`` / ``SequenceKey`` /
-      ``GetAttrKey`` / ``FlattenedIndexKey`` entries).
-
-  Returns:
-    The list of string key segments, e.g. ``["decoder", "layers_0", "kernel"]``.
-  """
-  parts: list[str] = []
-  for entry in path_tuple:
-    key = getattr(entry, "key", getattr(entry, "idx", getattr(entry, "name", entry)))
-    # Fold a layer/expert index (an int, or a numeric string after an Orbax
-    # round-trip) into the preceding segment: ["layers", 0] -> "layers_0".
-    if (isinstance(key, int) or (isinstance(key, str) and key.isdigit())) and parts:
-      parts[-1] = f"{parts[-1]}_{key}"
-    else:
-      parts.append(str(key))
-  # Drop the trailing ``value`` segment that NNX adds for each ``nnx.Variable``
-  # leaf (``...-kernel -> value``); MaxText-Linen param keys have no such wrapper.
-  if parts and parts[-1] == "value":
-    parts.pop()
-  return parts
-
-
 def extract_nnx_weights(weights_dict: dict) -> dict[str, np.ndarray]:
   """Extract weights from NNX checkpoint structure.
 
@@ -948,10 +903,13 @@ def extract_nnx_weights(weights_dict: dict) -> dict[str, np.ndarray]:
   result = {}
   leaves_with_paths = jax.tree_util.tree_leaves_with_path(weights_dict)
   for path_tuple, leaf_value in leaves_with_paths:
-    path_keys = param_key_parts_from_path(path_tuple)
+    path_keys = [k.key for k in path_tuple]
     # Skip NNX RNG state variables (not model weights)
     if "to_nnx__rngs" in path_keys or any(k.endswith("_rngs") for k in path_keys):
       continue
+    # Skip if this is the "value" key itself - we want the parent path
+    if path_keys[-1] == "value":
+      path_keys = path_keys[:-1]
     maxtext_param_key = "params-" + "-".join(path_keys)
     if not isinstance(leaf_value, (jax.Array, np.ndarray)):
       raise ValueError(f"Leaf value for {maxtext_param_key} is not an array. Type: {type(leaf_value)}.")
@@ -974,7 +932,8 @@ def extract_linen_weights(weights_dict: dict) -> dict[str, np.ndarray]:
   result = {}
   leaves_with_paths = jax.tree_util.tree_leaves_with_path(weights_dict)
   for path_tuple, leaf_value in leaves_with_paths:
-    path_keys = param_key_parts_from_path(path_tuple)
+    path_keys = [k.key for k in path_tuple]
+    # Construct maxtext_param_key from path_tuple
     maxtext_param_key = "params-" + "-".join(path_keys)
     if not isinstance(leaf_value, (jax.Array, np.ndarray)):
       raise ValueError(f"Leaf value for {maxtext_param_key} is not an array. Type: {type(leaf_value)}.")

src/maxtext/common/goodput.py

@@ -22,6 +22,7 @@
 import contextlib
 import jax
 from enum import Enum
+import pathwaysutils
 from maxtext.utils import max_logging
 from maxtext.common.gcloud_stub import goodput_modules
 
@@ -66,18 +67,31 @@ def maybe_monitor_goodput(config):
         enable_gcp_goodput_metrics=config.enable_gcp_goodput_metrics,
         enable_gcp_step_deviation_metrics=config.enable_gcp_step_deviation_metrics,
     )
-    goodput_monitor = monitoring.GoodputMonitor(
+    monitor_class = monitoring.GoodputMonitor
+
+    if config.elastic_enabled:
+      try:
+        from ml_goodput_measurement import monitoring_elastic
+        monitor_class = monitoring_elastic.ElasticGoodputMonitor
+      except ImportError:
+        max_logging.log("Elastic monitor failed!")
+        pass
+
+    kwargs = dict(
         job_name=config.run_name,
         logger_name=f"goodput_{config.run_name}",
         tensorboard_dir=config.tensorboard_dir,
         upload_interval=config.goodput_upload_interval_seconds,
         monitoring_enabled=True,
-        pathway_enabled=config.enable_pathways_goodput,
         include_badput_breakdown=True,
         include_step_deviation=config.monitor_step_time_deviation,
         step_deviation_interval_seconds=config.step_deviation_interval_seconds,
         gcp_options=gcp_options,
     )
+    if monitor_class == monitoring.GoodputMonitor:
+        kwargs["pathway_enabled"] = config.enable_pathways_goodput
+
+    goodput_monitor = monitor_class(**kwargs)
     goodput_monitor.start_goodput_uploader()
     max_logging.log("Started Goodput upload to Tensorboard & GCM in the background!")
     yield
@@ -121,8 +135,22 @@ def create_goodput_recorder(config):
     if config.enable_goodput_recording and jax.process_index() == 0:
       max_logging.log("[GOODPUT NO-OP] recorder skipped (decoupled stub).")
     return None
-  if config.enable_goodput_recording:
-    logger_name = f"goodput_{config.run_name}"
-    recorder = goodput.GoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)
-    return recorder
-  return None
+    
+  if not config.enable_goodput_recording:
+    return None
+
+  logger_name = f"goodput_{config.run_name}"
+
+  # Detect if we should use the elastic-aware recorder
+  if config.elastic_enabled:
+    try:
+      from ml_goodput_measurement import goodput_elastic
+      from maxtext.utils import elastic_utils
+      recorder = goodput_elastic.ElasticGoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)      
+      elastic_utils.record_slice_state(recorder)
+    except ImportError as e:
+      max_logging.log(f"Could not create elastic goodput recorder: {e}")
+    else:
+      return recorder
+
+  return goodput.GoodputRecorder(config.run_name, logger_name, jax.process_index() == 0)