Merge branch 'main' into htn/fix-jax-splash-soft-cap

Author: huytransformer

.github/CODEOWNERS

@@ -1,11 +1,11 @@
-* @abhinavclemson @gobbleturk @khatwanimohit @bvandermoon @vipannalla @RissyRan @richjames0 @gagika @shralex @SurbhiJainUSC @hengtaoguo @A9isha @aireenmei @NuojCheng @jiangjy1982 @suexu1025 @NicoGrande @jesselu-google @dipannita08 @igorts-git @darisoy
+* @abhinavclemson @gobbleturk @khatwanimohit @bvandermoon @vipannalla @RissyRan @richjames0 @gagika @shralex @SurbhiJainUSC @hengtaoguo @A9isha @aireenmei @NuojCheng @jiangjy1982 @suexu1025 @NicoGrande @jesselu-google @dipannita08 @igorts-git @darisoy @xibinliu
 
 # Model bring-up
 src/maxtext/assets @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande
 src/maxtext/configs/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng
 src/maxtext/checkpoint_conversion @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @hengtaoguo @gagika @shralex @richjames0 @NicoGrande
-src/maxtext/layers @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng
-src/maxtext/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng
+src/maxtext/layers @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng @xibinliu
+src/maxtext/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng @xibinliu
 
 # Features
 src/maxtext/experimental/rl @A9isha @khatwanimohit @xuefgu @gagika @richjames0 @shralex @NicoGrande @darisoy
@@ -32,7 +32,7 @@ src/dependencies/ @bvandermoon @SurbhiJainUSC @parambole @richjames0 @shralex @d
 docs/ @jacoguzo @bvandermoon @SurbhiJainUSC @richjames0 @shralex @gobbleturk @RissyRan @gagika @A9isha @jiangjy1982 @vipannalla @darisoy
 
 # Workflow files
-.github/workflows/ @gobbleturk @khatwanimohit @shralex @parambole @bvandermoon @richjames0 @darisoy
+.github/workflows/ @gobbleturk @khatwanimohit @shralex @parambole @bvandermoon @richjames0 @darisoy @xibinliu
 
 # Benchmarking/Recipes
 benchmarks/ @SujeethJinesh @bvandermoon @richjames0 @shralex @vipannalla @mitalisi @RissyRan @shauryagup @NuojCheng @gobbleturk @khatwanimohit @Obliviour @notabee @suexu1025 @darisoy

.github/workflows/ci_pipeline.yml

@@ -199,12 +199,31 @@ jobs:
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
 
+  setup-pathways-parameters:
+    name: Setup Pathways Parameters
+    runs-on: ubuntu-latest
+    outputs:
+      total_workers: ${{ steps.set-params.outputs.total_workers }}
+      worker_groups: ${{ steps.set-params.outputs.worker_groups }}
+    steps:
+      - id: set-params
+        name: Set Pathways Worker Parameters
+        run: |
+          # TPU worker constants
+          TPU_UNIT_TOTAL_WORKERS=2
+          TPU_UNIT_WORKER_GROUPS='[1, 2]'
+
+          echo "total_workers=${TPU_UNIT_TOTAL_WORKERS}" >> "$GITHUB_OUTPUT"
+          echo "worker_groups=${TPU_UNIT_WORKER_GROUPS}" >> "$GITHUB_OUTPUT"
+
   maxtext_tpu_pathways_unit_tests:
-    needs: build_and_upload_maxtext_package
+    needs: [build_and_upload_maxtext_package, setup-pathways-parameters]
     if: needs.analyze_code_changes.outputs.run_tests == 'true'
     uses: ./.github/workflows/run_pathways_tests.yml
     strategy:
         fail-fast: false
+        matrix:
+          group: ${{ fromJSON(needs.setup-pathways-parameters.outputs.worker_groups) }}
     with:
       device_type: tpu
       device_name: v6e-4
@@ -217,6 +236,8 @@ jobs:
       container_resource_option: "--privileged"
       is_scheduled_run: ${{ github.event_name == 'schedule' }}
       maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
+      total_workers: ${{ needs.setup-pathways-parameters.outputs.total_workers }}
+      worker_group: ${{ matrix.group }}
 
   maxtext_tpu_pathways_integration_tests:
     needs: build_and_upload_maxtext_package

.github/workflows/run_pathways_tests.yml

@@ -53,6 +53,15 @@ on:
       maxtext_sha:
         required: true
         type: string
+      total_workers:
+        required: false
+        type: string
+        default: '1'
+      worker_group:
+        required: false
+        type: string
+        default: '1'
+
 
 permissions:
   contents: read
@@ -95,6 +104,12 @@ jobs:
         run : gcloud storage cp gs://maxtext-test-assets/* tests/assets
       - name: Run Tests
         run: |
+          if [ "${{ inputs.total_workers }}" -gt 1 ]; then
+            .venv/bin/python3 -m pip install --quiet pytest-split
+            SPLIT_ARGS="--splits ${{ inputs.total_workers }} --group ${{ inputs.worker_group }}"
+          else
+            SPLIT_ARGS=""
+          fi
           if [ "${{ inputs.is_scheduled_run }}" = "true" ]; then
             FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }}"
           else
@@ -105,7 +120,7 @@ jobs:
           export MAXTEXT_TEST_ASSETS_ROOT=$(pwd)/tests/assets
           export MAXTEXT_PKG_DIR=$(pwd)/src/maxtext
           # TODO(b/454659463): Enable test_default_hlo_match after volume mount is supported.
-          .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" -k "not AotHloIdenticalTest and not CompileThenLoad" --durations=0
+          .venv/bin/python3 -m pytest ${{ inputs.pytest_addopts }} -v -m "${FINAL_PYTEST_MARKER}" -k "not AotHloIdenticalTest and not CompileThenLoad" --durations=0 ${SPLIT_ARGS}
         env:
           PYTHONPATH: "${{ github.workspace }}/src"
     services:

.github/workflows/run_tests_against_package.yml

@@ -176,8 +176,13 @@ jobs:
             fi
           fi
           if [ "${INPUTS_TOTAL_WORKERS}" -gt 1 ]; then
-            $PYTHON_EXE -m pip install --quiet pytest-split pytest-xdist
-            SPLIT_ARGS="--splits ${INPUTS_TOTAL_WORKERS} --group ${INPUTS_WORKER_GROUP} -n auto"
+            $PYTHON_EXE -m pip install --quiet pytest-split
+            if [ "${INPUTS_DEVICE_TYPE}" = "tpu" ]; then
+              SPLIT_ARGS="--splits ${INPUTS_TOTAL_WORKERS} --group ${INPUTS_WORKER_GROUP}"
+            else
+              $PYTHON_EXE -m pip install --quiet pytest-xdist
+              SPLIT_ARGS="--splits ${INPUTS_TOTAL_WORKERS} --group ${INPUTS_WORKER_GROUP} -n auto"
+            fi
           else
             SPLIT_ARGS=""
           fi

.github/workflows/run_tests_coordinator.yml

@@ -61,12 +61,48 @@ permissions:
   contents: read
 
 jobs:
+  setup-parameters:
+    name: Setup Parameters
+    runs-on: ubuntu-latest
+    outputs:
+      worker_groups: ${{ steps.set-params.outputs.worker_groups }}
+      total_workers: ${{ steps.set-params.outputs.total_workers }}
+    steps:
+      - id: set-params
+        name: Set Worker Parameters
+        run: |
+          # CPU worker constants
+          CPU_UNIT_TOTAL_WORKERS=4
+          CPU_UNIT_WORKER_GROUPS='[1, 2, 3, 4]'
+
+          # TPU worker constants
+          TPU_UNIT_TOTAL_WORKERS=2
+          TPU_UNIT_WORKER_GROUPS='[1, 2]'
+
+          # Default fallback constants
+          DEFAULT_TOTAL_WORKERS=1
+          DEFAULT_WORKER_GROUPS='[1]'
+
+          FLAVOR="${{ inputs.flavor }}"
+
+          if [[ "$FLAVOR" == "cpu-unit" || "$FLAVOR" == "cpu-post-training-unit" ]]; then
+            echo "worker_groups=${CPU_UNIT_WORKER_GROUPS}" >> "$GITHUB_OUTPUT"
+            echo "total_workers=${CPU_UNIT_TOTAL_WORKERS}" >> "$GITHUB_OUTPUT"
+          elif [[ "$FLAVOR" == "tpu-unit" ]]; then
+            echo "worker_groups=${TPU_UNIT_WORKER_GROUPS}" >> "$GITHUB_OUTPUT"
+            echo "total_workers=${TPU_UNIT_TOTAL_WORKERS}" >> "$GITHUB_OUTPUT"
+          else
+            echo "worker_groups=${DEFAULT_WORKER_GROUPS}" >> "$GITHUB_OUTPUT"
+            echo "total_workers=${DEFAULT_TOTAL_WORKERS}" >> "$GITHUB_OUTPUT"
+          fi
+
   execute-test-package:
+    needs: setup-parameters
     name: ${{ inputs.flavor }}
     strategy:
       fail-fast: false
       matrix:
-        worker_group: ${{ fromJSON(contains(inputs.flavor, 'cpu-unit') && '[1, 2, 3, 4]' || '[1]') }}
+        worker_group: ${{ fromJSON(needs.setup-parameters.outputs.worker_groups) }}
 
     uses: ./.github/workflows/run_tests_against_package.yml
     with:
@@ -158,6 +194,6 @@ jobs:
       is_scheduled_run: ${{ inputs.is_scheduled_run }}
       maxtext_installed: ${{ inputs.maxtext_installed }}
       worker_group: ${{ matrix.worker_group }}
-      total_workers: ${{ contains(inputs.flavor, 'cpu-unit') && 4 || 1 }}
+      total_workers: ${{ fromJSON(needs.setup-parameters.outputs.total_workers) }}
       maxtext_sha: ${{ inputs.maxtext_sha }}
       is_update_hlo: ${{ inputs.is_update_hlo }}

docs/guides/lora_model_bringup.md

@@ -30,12 +30,12 @@ To enable LoRA support for a new model, follow these two simple steps:
 
 The target model architecture must already be implemented and supported as a base model in MaxText.
 
-- The JAX/NNX model definition should be located under `src/maxtext/models/` (e.g., \[gemma3.py\](../../src/maxtext/models/gemma3.py)).
+- The JAX/NNX model definition should be located under `src/maxtext/models/` (e.g., [gemma3.py](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/models/gemma3.py)).
 - The model configurations must be registered and runnable for baseline pre-training or full fine-tuning.
 
 ### Step 1.2: Define Trainable LoRA Target Modules
 
-Add a recommended target pattern for your model architecture prefix in \[src/maxtext/configs/post_train/lora_module_path.yml\](../../src/maxtext/configs/post_train/lora_module_path.yml):
+Add a recommended target pattern for your model architecture prefix in [src/maxtext/configs/post_train/lora_module_path.yml](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/configs/post_train/lora_module_path.yml):
 
 ```yaml
 your_model_prefix: "decoder/layers/.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))"
@@ -69,7 +69,7 @@ If you want to perform decoding or run high-performance serving on your adapted
 To add weight mapping for vLLM decode:
 
 1. **Create a Weight Mapping Config**:
-   Create a new file in \[src/maxtext/integration/tunix/weight_mapping/\](../../src/maxtext/integration/tunix/weight_mapping/) (e.g., `your_model.py`) defining a mapping dataclass. You can refer to \[gemma3.py\](../../src/maxtext/integration/tunix/weight_mapping/gemma3.py) or \[llama3.py\](../../src/maxtext/integration/tunix/weight_mapping/llama3.py) as templates.
+   Create a new file in [src/maxtext/integration/tunix/weight_mapping/](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/integration/tunix/weight_mapping/) (e.g., `your_model.py`) defining a mapping dataclass. You can refer to [llama3.py](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/integration/tunix/weight_mapping/llama3.py) as a template.
 
    Your class should specify:
 
@@ -78,7 +78,7 @@ To add weight mapping for vLLM decode:
    - `lora_to_hf_mappings()`: Custom mapping for LoRA weights if they require different handling.
 
 2. **Register the Mapping**:
-   Register your new class in \[src/maxtext/integration/tunix/weight_mapping/__init__.py\](../../src/maxtext/integration/tunix/weight_mapping/__init__.py) inside the `StandaloneVllmWeightMapping` class:
+   Register your new class in [src/maxtext/integration/tunix/weight_mapping/__init__.py](https://github.com/AI-Hypercomputer/maxtext/blob/main/src/maxtext/integration/tunix/weight_mapping/__init__.py) inside the `StandaloneVllmWeightMapping` class:
 
    ```python
    # Inside StandaloneVllmWeightMapping

docs/reference/core_concepts/moe_configuration.md

@@ -99,6 +99,8 @@ Dropping:
 
 `mlp_bias`: If enabled, add learnable bias terms for MLP matmul. Originally implemented to support the GPT-OSS model architecture.
 
+`prefuse_moe_weights`: If enabled alongside `sparse_matmul=True`, fuses the two FFN1 grouped GEMMs (wi\_0 and wi\_1) into a single grouped GEMM call. Expert weights are stored in a concatenated `(num_experts, embed_dim, 2 * mlp_dim)` shape, so input activations are loaded from HBM once per forward pass instead of twice. Backend-agnostic (works with Megablox, JAX Ragged Dot, and Tokamax). When used with `attention=vllm_rpa`, the fused weight tensor is passed directly to the vLLM-TPU serving kernel without splitting.
+
 `use_batch_split_schedule` (experimental): If enabled, split batch into micro-batches to hide communications that yields performance benefits.
 
 ## 2. Sharding

src/maxtext/checkpoint_conversion/to_maxtext.py

@@ -67,7 +67,8 @@
 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.tensor_handling import apply_hook_fns
+from maxtext.checkpoint_conversion.utils.utils import MemoryMonitorTqdm, 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.inference.inference_utils import str2bool
 from maxtext.layers import quantizations
 from maxtext.models import models
@@ -319,22 +320,21 @@ def get_maxtext_model_info(config):
   # Get abstract model structure (name, shape) without materializing the weights to save memory
   abstract_params_tree = maxtext_utils.get_abstract_param(maxtext_model_flax, config)["params"]
 
-  abstract_params_flat, _ = jax.tree_util.tree_flatten_with_path(abstract_params_tree)
-  # Standardize abstract tree for later unflattening
-  abstract_params_tree = jax.tree.map(
-      lambda _: 0,
+  abstract_params_flat, abstract_params_treedef = jax.tree_util.tree_flatten_with_path(
       abstract_params_tree,
       is_leaf=lambda x: isinstance(x, nn.LogicallyPartitioned),
   )
-  abstract_params_treedef = jax.tree_util.tree_structure(abstract_params_tree)
 
   max_logging.log("MaxText abstract model and state initialized.")
 
   # 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))
-    mt_target_shape = abstract_leaf_value.shape
+    if isinstance(abstract_leaf_value, nn.LogicallyPartitioned):
+      mt_target_shape = abstract_leaf_value.value.shape
+    else:
+      mt_target_shape = abstract_leaf_value.shape
     maxtext_abstract_dict[mt_param_key] = (mt_target_idx, mt_target_shape)
 
   return maxtext_abstract_dict, abstract_params_treedef

src/maxtext/checkpoint_conversion/utils/hf_model_configs.py

@@ -1642,6 +1642,62 @@ def __init__(self, **kwargs):
 }
 qwen3_vl_4b_config = PTConfig(**qwen3_vl_4b_dict)
 
+qwen3_vl_2b_dict = {
+    "architectures": ["Qwen3VLForConditionalGeneration"],
+    "image_token_id": 151655,
+    "model_type": "qwen3_vl",
+    "text_config": {
+        "attention_bias": False,
+        "attention_dropout": 0.0,
+        "bos_token_id": 151643,
+        "dtype": "bfloat16",
+        "eos_token_id": 151645,
+        "head_dim": 128,
+        "hidden_act": "silu",
+        "hidden_size": 2048,
+        "initializer_range": 0.02,
+        "intermediate_size": 6144,
+        "max_position_embeddings": 262144,
+        "model_type": "qwen3_vl_text",
+        "num_attention_heads": 16,
+        "num_hidden_layers": 28,
+        "num_key_value_heads": 8,
+        "pad_token_id": None,
+        "rms_norm_eps": 1e-06,
+        "rope_parameters": {
+            "mrope_interleaved": True,
+            "mrope_section": [24, 20, 20],
+            "rope_theta": 5000000,
+            "rope_type": "default",
+        },
+        "tie_word_embeddings": True,
+        "use_cache": True,
+        "vocab_size": 151936,
+    },
+    "tie_word_embeddings": True,
+    "transformers_version": "4.57.0.dev0",
+    "video_token_id": 151656,
+    "vision_config": {
+        "deepstack_visual_indexes": [5, 11, 17],
+        "depth": 24,
+        "hidden_act": "gelu_pytorch_tanh",
+        "hidden_size": 1024,
+        "in_channels": 3,
+        "initializer_range": 0.02,
+        "intermediate_size": 4096,
+        "model_type": "qwen3_vl_vision",
+        "num_heads": 16,
+        "num_position_embeddings": 2304,
+        "out_hidden_size": 2048,
+        "patch_size": 16,
+        "spatial_merge_size": 2,
+        "temporal_patch_size": 2,
+    },
+    "vision_end_token_id": 151653,
+    "vision_start_token_id": 151652,
+}
+qwen3_vl_2b_config = PTConfig(**qwen3_vl_2b_dict)
+
 
 # {maxtext model name: hf model config}
 HF_MODEL_CONFIGS = {
@@ -1669,6 +1725,7 @@ def __init__(self, **kwargs):
     "qwen3-14b": qwen3_14b_config,
     "qwen3-14b-base": qwen3_14b_config,
     "qwen3-32b": qwen3_32b_config,
+    "qwen3-vl-2b": qwen3_vl_2b_config,
     "qwen3-vl-4b": qwen3_vl_4b_config,
     "llama3.1-8b": llama31_8b_config,
     "llama3.1-8b-Instruct": llama31_8b_config,

src/maxtext/checkpoint_conversion/utils/hf_shape.py

@@ -1182,6 +1182,7 @@ def QWEN3_VL_HF_WEIGHTS_TO_SHAPE(config):
     "qwen3-8b": QWEN_HF_WEIGHTS_TO_SHAPE,
     "qwen3-14b": QWEN_HF_WEIGHTS_TO_SHAPE,
     "qwen3-32b": QWEN_HF_WEIGHTS_TO_SHAPE,
+    "qwen3-vl-2b": QWEN3_VL_HF_WEIGHTS_TO_SHAPE,
     "qwen3-vl-4b": QWEN3_VL_HF_WEIGHTS_TO_SHAPE,
     "llama3.1-8b": LLAMA31_HF_WEIGHTS_TO_SHAPE,
     "llama3.1-8b-Instruct": LLAMA31_HF_WEIGHTS_TO_SHAPE,