From cdd82bde7c917426413bd18559885cbf8cafd823 Mon Sep 17 00:00:00 2001 From: Surbhi Jain Date: Thu, 13 Aug 2026 19:02:38 +0000 Subject: [PATCH] Fix CI notebook execution and add SFT training step verification --- .github/workflows/ci_pipeline.yml | 4 +- .github/workflows/run_ci_tests.yml | 2 +- .github/workflows/run_jupyter_notebooks.yml | 119 ++++++++---- .github/workflows/run_tests_coordinator.yml | 18 +- src/maxtext/examples/lora_llama3_demo.ipynb | 14 +- .../examples/sft_llama3_demo_tpu.ipynb | 14 +- .../examples/sft_multimodal_gemma3_demo.ipynb | 2 +- src/maxtext/examples/sft_qwen3_demo.ipynb | 14 +- .../integration/vllm/maxtext_vllm_rollout.py | 1 - src/maxtext/layers/nnx_scan.py | 44 ++++- src/maxtext/layers/quantizations.py | 18 ++ .../trainers/post_train/sft/train_sft.py | 38 ++-- src/maxtext/utils/lora_utils.py | 30 +++ src/maxtext/utils/maxtext_utils.py | 8 +- src/maxtext/utils/sharding.py | 12 ++ tests/integration/lora_e2e_nnx_test.py | 94 +--------- .../integration/lora_e2e_nnx_test.py | 175 ++++++++++++++++++ tests/unit/sharding_nnx_test.py | 32 ++-- 18 files changed, 454 insertions(+), 185 deletions(-) create mode 100644 tests/post_training/integration/lora_e2e_nnx_test.py diff --git a/.github/workflows/ci_pipeline.yml b/.github/workflows/ci_pipeline.yml index e854a023d6..2b69c20ac9 100644 --- a/.github/workflows/ci_pipeline.yml +++ b/.github/workflows/ci_pipeline.yml @@ -168,7 +168,7 @@ jobs: strategy: fail-fast: false matrix: - flavor: [tpu-post-training-unit] + flavor: [tpu-post-training-unit, tpu-post-training-integration] with: flavor: ${{ matrix.flavor }} base_image: maxtext-unit-test-tpu:py312 @@ -241,7 +241,7 @@ jobs: strategy: fail-fast: false matrix: - flavor: [cpu-post-training-unit] + flavor: [cpu-post-training-unit, cpu-post-training-integration] with: flavor: ${{ matrix.flavor }} base_image: maxtext-unit-test-tpu:py312 diff --git a/.github/workflows/run_ci_tests.yml b/.github/workflows/run_ci_tests.yml index b51ccf6ee3..e7d98ae956 100644 --- a/.github/workflows/run_ci_tests.yml +++ b/.github/workflows/run_ci_tests.yml @@ -71,7 +71,7 @@ jobs: flavor: >- ${{ fromJSON('{ "gpu-pre-training": ["gpu-unit", "gpu-integration"], - "tpu-post-training": ["tpu-post-training-unit", "tpu-post-training-integration", "cpu-post-training-unit"], + "tpu-post-training": ["tpu-post-training-unit", "tpu-post-training-integration", "cpu-post-training-unit", "cpu-post-training-integration"], "tpu-pre-training": ["tpu-unit", "tpu-integration", "cpu-unit", "cpu-integration"] }')[format('{0}-{1}', inputs.device, inputs.workflow)] }} uses: ./.github/workflows/run_tests_coordinator.yml diff --git a/.github/workflows/run_jupyter_notebooks.yml b/.github/workflows/run_jupyter_notebooks.yml index d102837df3..706baea8bc 100644 --- a/.github/workflows/run_jupyter_notebooks.yml +++ b/.github/workflows/run_jupyter_notebooks.yml @@ -34,11 +34,6 @@ on: maxtext_sha: required: false type: string - # Flag to skip source checkout and wheel installation - maxtext_installed: - required: false - type: boolean - default: false secrets: HF_TOKEN: required: true @@ -46,8 +41,70 @@ on: permissions: contents: read jobs: + discover_notebooks: + name: Discover Notebooks to Run + runs-on: ubuntu-latest + outputs: + notebooks: ${{ steps.list.outputs.notebooks }} + count: ${{ steps.list.outputs.count }} + steps: + - name: Checkout MaxText + uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5.1.0 + with: + ref: ${{ inputs.maxtext_sha }} + fetch-depth: 0 + persist-credentials: false + - name: Determine Notebook List + id: list + shell: bash + run: | + MAXTEXT_NOTEBOOKS_ROOT="src/maxtext/examples" + SKIPPED=("sft_llama3_demo_gpu.ipynb" "maxtext_with_gepa.ipynb" "demo_decoding.ipynb" "dpo_qwen3_demo.ipynb") + + is_skipped() { + local name="$1" + for s in "${SKIPPED[@]}"; do + [[ "$name" == "$s" ]] && return 0 + done + return 1 + } + + SELECTED=() + if [ "${GITHUB_EVENT_NAME}" == "pull_request" ] && [ -n "${GITHUB_BASE_REF}" ]; then + git fetch origin "${GITHUB_BASE_REF}" --depth=50 2>/dev/null || true + while IFS= read -r file; do + if [[ -n "$file" && -f "$file" ]]; then + name=$(basename "$file") + if ! is_skipped "$name"; then + SELECTED+=("$name") + fi + fi + done < <(git diff --name-only "origin/${GITHUB_BASE_REF}...HEAD" -- "${MAXTEXT_NOTEBOOKS_ROOT}"/*.ipynb 2>/dev/null || true) + fi + + # If not a PR or no specific notebooks were modified (e.g. scheduled run or workflow file modified), run all active notebooks + if [ ${#SELECTED[@]} -eq 0 ]; then + for nb in "$MAXTEXT_NOTEBOOKS_ROOT"/*.ipynb; do + name=$(basename "$nb") + if ! is_skipped "$name" && [[ -f "$nb" ]]; then + SELECTED+=("$name") + fi + done + fi + + JSON_ARRAY=$(jq -nc '$ARGS.positional' --args "${SELECTED[@]}") + echo "Discovered notebooks to run: $JSON_ARRAY" + echo "notebooks=$JSON_ARRAY" >> "$GITHUB_OUTPUT" + echo "count=${#SELECTED[@]}" >> "$GITHUB_OUTPUT" + run: - name: Execute Notebooks + name: Execute ${{ matrix.notebook }} + needs: [discover_notebooks] + if: needs.discover_notebooks.outputs.count > 0 + strategy: + fail-fast: false + matrix: + notebook: ${{ fromJson(needs.discover_notebooks.outputs.notebooks) }} runs-on: ${{ inputs.cloud_runner != '' && inputs.cloud_runner || fromJson(format('["self-hosted", "{0}", "{1}"]', inputs.device_type, inputs.device_name)) }} container: image: gcr.io/tpu-prod-env-multipod/${{ inputs.base_image }} # zizmor: ignore[unpinned-images] @@ -56,18 +113,15 @@ jobs: UV_TORCH_BACKEND: "cpu" steps: - name: Checkout MaxText - if: ${{ !inputs.maxtext_installed }} uses: actions/checkout@fbc6f3992d24b796d5a048ff273f7fcc4a7b6c09 # v5.1.0 with: ref: ${{ inputs.maxtext_sha }} persist-credentials: false - name: Download the MaxText wheel - if: ${{ !inputs.maxtext_installed }} uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdf5c8093 # v4.3.0 with: name: maxtext-wheel - name: Install MaxText and Dependencies - if: ${{ !inputs.maxtext_installed }} shell: bash run: | # 1. Create virtual environment @@ -89,25 +143,18 @@ jobs: install_tpu_post_train_extra_deps python3 -m pip freeze - - name: Run Post-Training Notebooks + - name: Run Post-Training Notebook shell: bash env: HF_TOKEN: ${{ secrets.HF_TOKEN }} - MAXTEXT_INSTALLED: ${{ inputs.maxtext_installed }} + NOTEBOOK_NAME: ${{ matrix.notebook }} # TODO: Fix evaluation in sft_qwen3_demo.ipynb and remove this env variable RUN_EVALUATION: "false" run: | - if [ "${MAXTEXT_INSTALLED}" == "true" ]; then - # Move to the directory where code is baked into the image. See the Dockerfile. - # This is necessary because GHA sets an empty workspace by default. - cd /deps - PYTHON_EXE="python3" - PAPERMILL_EXE="papermill" - else - PYTHON_EXE=".venv/bin/python3" - PAPERMILL_EXE=".venv/bin/papermill" - source .venv/bin/activate - fi + PYTHON_EXE=".venv/bin/python3" + PAPERMILL_EXE=".venv/bin/papermill" + source .venv/bin/activate + export PYTHONPATH="${PWD}/src${PYTHONPATH:+:${PYTHONPATH}}" MAXTEXT_REPO_ROOT=$(pwd) @@ -121,28 +168,22 @@ jobs: # Run Hugging Face authentication hf auth login --token "$HF_TOKEN" - for notebook in "$MAXTEXT_NOTEBOOKS_ROOT"/*.ipynb; do - filename=$(basename "$notebook") - if [[ "$filename" == "sft_llama3_demo_gpu.ipynb" || "$filename" == "maxtext_with_gepa.ipynb" || "$filename" == "demo_decoding.ipynb" || "$filename" == "dpo_qwen3_demo.ipynb" ]]; then - echo "Skipping $filename" - continue - fi - output_name="${filename%.ipynb}_output.ipynb" + notebook="$MAXTEXT_NOTEBOOKS_ROOT/$NOTEBOOK_NAME" + output_name="${NOTEBOOK_NAME%.ipynb}_output.ipynb" - echo "------------------------------------------------------" - echo "Running $filename ..." - echo "------------------------------------------------------" + echo "------------------------------------------------------" + echo "Running $NOTEBOOK_NAME ..." + echo "------------------------------------------------------" - $PAPERMILL_EXE "$notebook" "$output_name" -k maxtext_venv + $PAPERMILL_EXE "$notebook" "$output_name" -k maxtext_venv - # Clean up any checkpoint directories created by the notebook to avoid filling up disk space - echo "Post-notebook disk cleanup for $filename ..." - rm -rf "$MAXTEXT_PKG_DIR"/*_output - rm -rf "$HOME/.cache/huggingface/hub" - done + # Clean up any checkpoint directories created by the notebook to avoid filling up disk space + echo "Post-notebook disk cleanup for $NOTEBOOK_NAME ..." + rm -rf "$MAXTEXT_PKG_DIR"/*_output + rm -rf "$HOME/.cache/huggingface/hub" - name: Upload Outputs if: always() uses: actions/upload-artifact@5d5d22a31266ced268874388b861e4b58bb5c2f3 # v4.3.1 with: - name: notebook-outputs-${{ inputs.device_name }} + name: notebook-outputs-${{ matrix.notebook }}-${{ inputs.device_name }} path: ./*_output.ipynb diff --git a/.github/workflows/run_tests_coordinator.yml b/.github/workflows/run_tests_coordinator.yml index 96be70f92f..bc6fcf4fb2 100644 --- a/.github/workflows/run_tests_coordinator.yml +++ b/.github/workflows/run_tests_coordinator.yml @@ -119,7 +119,8 @@ jobs: "gpu-integration": "cuda12", "cpu-unit": "cpu", "cpu-integration": "cpu", - "cpu-post-training-unit": "cpu" + "cpu-post-training-unit": "cpu", + "cpu-post-training-integration": "cpu" }')[inputs.flavor] }} device_name: >- @@ -136,7 +137,8 @@ jobs: "gpu-integration": "a100-40gb-4", "cpu-unit": "X64", "cpu-integration": "X64", - "cpu-post-training-unit": "X64" + "cpu-post-training-unit": "X64", + "cpu-post-training-integration": "X64" }')[inputs.flavor] }} cloud_runner: >- @@ -153,7 +155,8 @@ jobs: "gpu-integration": "linux-x86-a2-48-a100-4gpu", "cpu-unit": "linux-x86-n2-32", "cpu-integration": "linux-x86-n2-32", - "cpu-post-training-unit": "linux-x86-n2-32" + "cpu-post-training-unit": "linux-x86-n2-32", + "cpu-post-training-integration": "linux-x86-n2-32" }')[inputs.flavor] }} # Pytest Marker Mapping pytest_marker: >- @@ -170,7 +173,8 @@ jobs: "gpu-integration": "not cpu_only and not tpu_only and integration_test and not post_training", "cpu-unit": "cpu_only and not post_training and not integration_test", "cpu-integration": "cpu_only and not post_training and integration_test", - "cpu-post-training-unit": "cpu_only and post_training" + "cpu-post-training-unit": "cpu_only and post_training", + "cpu-post-training-integration": "cpu_only and post_training and integration_test" }')[inputs.flavor] }} pytest_addopts: >- @@ -187,7 +191,8 @@ jobs: "gpu-integration": "", "cpu-unit": "", "cpu-integration": "", - "cpu-post-training-unit": "tests/post_training/unit tests/unit" + "cpu-post-training-unit": "tests/post_training/unit tests/unit", + "cpu-post-training-integration": "tests/post_training/integration" }')[inputs.flavor] }} pytest_extra_args: >- @@ -204,7 +209,8 @@ jobs: "gpu-integration": "--ignore=tests/post_training", "cpu-unit": "--ignore=tests/post_training", "cpu-integration": "--ignore=tests/post_training", - "cpu-post-training-unit": "" + "cpu-post-training-unit": "", + "cpu-post-training-integration": "" }')[inputs.flavor] }} # Resource Scaling diff --git a/src/maxtext/examples/lora_llama3_demo.ipynb b/src/maxtext/examples/lora_llama3_demo.ipynb index 99ce8481f3..280b48bd9f 100644 --- a/src/maxtext/examples/lora_llama3_demo.ipynb +++ b/src/maxtext/examples/lora_llama3_demo.ipynb @@ -140,7 +140,11 @@ "from flax import nnx\n", "from etils import epath\n", "\n", - "print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")" + "print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")\n", + "\n", + "from absl import flags\n", + "if not flags.FLAGS.is_parsed():\n", + " flags.FLAGS.mark_as_parsed()" ] }, { @@ -301,7 +305,7 @@ " f\"train_split={TRAIN_DATA_SPLIT}\",\n", " f\"hf_data_dir={HF_DATA_DIR}\",\n", " f\"train_data_columns={TRAIN_DATA_COLUMNS}\",\n", - " \"steps=200\",\n", + " \"steps=5\",\n", " \"per_device_batch_size=1\",\n", " \"max_target_length=512\",\n", " \"learning_rate=5e-5\", \n", @@ -465,6 +469,12 @@ "source": [ "print(\"Starting LoRA SFT Training...\")\n", "trainer = train_sft.train_model(config, trainer, mesh)\n", + "# Verify that the expected number of steps actually executed\n", + "if trainer.train_steps < config.steps:\n", + " raise RuntimeError(\n", + " f\"Training ended prematurely! Expected {config.steps} steps, \"\n", + " f\"but only completed {trainer.train_steps} steps.\"\n", + " )\n", "print(\"LoRA SFT Training Complete!\")" ] }, diff --git a/src/maxtext/examples/sft_llama3_demo_tpu.ipynb b/src/maxtext/examples/sft_llama3_demo_tpu.ipynb index d4fba92eed..216769ac48 100644 --- a/src/maxtext/examples/sft_llama3_demo_tpu.ipynb +++ b/src/maxtext/examples/sft_llama3_demo_tpu.ipynb @@ -127,7 +127,11 @@ "from etils import epath\n", "\n", "\n", - "print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")" + "print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")\n", + "\n", + "from absl import flags\n", + "if not flags.FLAGS.is_parsed():\n", + " flags.FLAGS.mark_as_parsed()" ] }, { @@ -253,7 +257,7 @@ " f\"{MAXTEXT_PKG_DIR}/configs/post_train/sft.yml\",\n", " f\"load_parameters_path={MODEL_CHECKPOINT_PATH}\",\n", " f\"model_name={MODEL_NAME}\",\n", - " \"steps=100\",\n", + " \"steps=5\",\n", " \"per_device_batch_size=1\",\n", " \"max_target_length=1024\",\n", " \"learning_rate=2.0e-5\",\n", @@ -296,6 +300,12 @@ "\n", "try:\n", " trainer, mesh = train_sft.train(config)\n", + " # Verify that the expected number of steps actually executed\n", + " if trainer.train_steps < config.steps:\n", + " raise RuntimeError(\n", + " f\"Training ended prematurely! Expected {config.steps} steps, \"\n", + " f\"but only completed {trainer.train_steps} steps.\"\n", + " )\n", " print(\"\\n\" + \"=\" * 60)\n", " print(\"✅ Training Completed Successfully!\")\n", " print(\"=\" * 60)\n", diff --git a/src/maxtext/examples/sft_multimodal_gemma3_demo.ipynb b/src/maxtext/examples/sft_multimodal_gemma3_demo.ipynb index b988649853..098e2f8d74 100644 --- a/src/maxtext/examples/sft_multimodal_gemma3_demo.ipynb +++ b/src/maxtext/examples/sft_multimodal_gemma3_demo.ipynb @@ -236,7 +236,7 @@ " f\"load_parameters_path={MODEL_CHECKPOINT_PATH}\",\n", " f\"model_name={MODEL_NAME}\",\n", " f\"tokenizer_path={TOKENIZER_NAME}\",\n", - " \"steps=10\",\n", + " \"steps=5\",\n", " \"attention=dot_product\",\n", " \"per_device_batch_size=1\",\n", " \"max_prefill_predict_length=1024\",\n", diff --git a/src/maxtext/examples/sft_qwen3_demo.ipynb b/src/maxtext/examples/sft_qwen3_demo.ipynb index e6618b8f49..0e0e838951 100644 --- a/src/maxtext/examples/sft_qwen3_demo.ipynb +++ b/src/maxtext/examples/sft_qwen3_demo.ipynb @@ -144,7 +144,11 @@ "from flax import nnx\n", "from etils import epath\n", "\n", - "print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")" + "print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")\n", + "\n", + "from absl import flags\n", + "if not flags.FLAGS.is_parsed():\n", + " flags.FLAGS.mark_as_parsed()" ] }, { @@ -315,7 +319,7 @@ " f\"train_split={TRAIN_DATA_SPLIT}\",\n", " f\"hf_data_dir={HF_DATA_DIR}\",\n", " f\"train_data_columns={TRAIN_DATA_COLUMNS}\",\n", - " \"steps=500\",\n", + " \"steps=5\",\n", " \"per_device_batch_size=1\",\n", " \"max_target_length=1024\",\n", " \"learning_rate=3e-6\",\n", @@ -488,6 +492,12 @@ "source": [ "print(\"Starting SFT Training...\")\n", "trainer = train_sft.train_model(config, trainer, mesh)\n", + "# Verify that the expected number of steps actually executed\n", + "if trainer.train_steps < config.steps:\n", + " raise RuntimeError(\n", + " f\"Training ended prematurely! Expected {config.steps} steps, \"\n", + " f\"but only completed {trainer.train_steps} steps.\"\n", + " )\n", "print(\"SFT Training Complete!\")" ] }, diff --git a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py index 636eee4d56..1845dbd8c0 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py @@ -291,7 +291,6 @@ def __init__( engine_kwargs = { "max_model_len": cache_config_or_size, "model": rollout_config.rollout_vllm_model_version, - "swap_space": getattr(rollout_config, "rollout_vllm_swap_space_size_gb", maxtext_config.swap_space_vllm_gb), # Async scheduling causes KeyError in dp_scheduler on slow models # (30B+) where inference latency exceeds the scheduler's window. "async_scheduling": rollout_config.rollout_vllm_async_scheduling, diff --git a/src/maxtext/layers/nnx_scan.py b/src/maxtext/layers/nnx_scan.py index e43a4d47bf..1198ef172e 100644 --- a/src/maxtext/layers/nnx_scan.py +++ b/src/maxtext/layers/nnx_scan.py @@ -126,14 +126,56 @@ def apply_scanned_layers( if param_scan_axis != 0: params = jax.tree.map(lambda x: jnp.moveaxis(x, param_scan_axis, 0), params) + def _ensure_stacked(x): + if hasattr(x, "ndim") and x.ndim == 0: + return jnp.broadcast_to(x, (length,)) + return x + + state = jax.tree.map(_ensure_stacked, state) + + def _strip_scan_metadata(leaf): + if hasattr(leaf, "replace") and hasattr(leaf, "value"): # pylint: disable=too-many-nested-blocks + replace_kwargs = {} + if hasattr(leaf, "get_metadata"): + replace_kwargs.update(leaf.get_metadata()) + + replace_kwargs.pop(nnx.PARTITION_NAME, None) + replace_kwargs.pop("param_scan_axis", None) + + val = getattr(leaf, "value", None) + val_ndim = getattr(val, "ndim", None) + + for key in ["sharding", "out_sharding", "kernel_axes", "sharding_names"]: + value = getattr(leaf, key, None) + if value is None and key in replace_kwargs: + value = replace_kwargs[key] + if value is not None: + if isinstance(value, str): + value = (value,) + if isinstance(value, tuple): + if val_ndim is not None and len(value) > val_ndim: + filtered = tuple( + axis for axis in value if axis not in ("local_layers", "layers", "scanned_blocks", "decoder_layers") + ) + if len(filtered) > val_ndim: + filtered = filtered[:val_ndim] + replace_kwargs[key] = filtered + return leaf.replace(**replace_kwargs) + return leaf + def scan_body(current_carry, scanned_state): current_params, current_state = scanned_state + current_params = jax.tree.map( + _strip_scan_metadata, + current_params, + is_leaf=lambda x: hasattr(x, "replace") and hasattr(x, "value"), + ) current_layer = nnx.merge(layer_graphdef, current_params, current_state) next_carry = apply_fn(current_layer, current_carry) return next_carry, nnx.state(current_layer) scan_fn = jax.checkpoint(scan_body, policy=remat_policy, prevent_cse=prevent_cse) if remat else scan_body - final_carry, scanned_state = jax.lax.scan(scan_fn, carry, (params, state), unroll=unroll) + final_carry, scanned_state = jax.lax.scan(scan_fn, carry, (params, state), length=length, unroll=unroll) if param_scan_axis != 0: scanned_params, scanned_other = scanned_state.split(nnx.Param, ...) diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index a5dade3e19..a275a0afa8 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -46,6 +46,21 @@ from qwix._src.utils import flax_util except ImportError: from qwix._src import flax_util # pytype: disable=import-error + +try: + _orig_find_param = flax_util.find_param + + def _safe_find_param(x, ptq_array_type=None): + try: + return _orig_find_param(x, ptq_array_type) + except AttributeError as e: + if "shape" in str(e): + return None + raise + + flax_util.find_param = _safe_find_param +except (NameError, AttributeError): + pass from maxtext.layers import nnx_wrappers from maxtext.configs.types import TeCommGemmOverlapPolicy @@ -888,6 +903,9 @@ def maybe_quantize_model(model, config): nnx.pop(model, nnx.Intermediate) else: model = qwix.quantize_model(model, quantization_provider) + for _, val in nnx.graph.iter_graph(model): + if hasattr(val, "__dict__") and "qwix_rngs" in val.__dict__: + del val.qwix_rngs return model diff --git a/src/maxtext/trainers/post_train/sft/train_sft.py b/src/maxtext/trainers/post_train/sft/train_sft.py index f46d6f6141..fb2926e67f 100644 --- a/src/maxtext/trainers/post_train/sft/train_sft.py +++ b/src/maxtext/trainers/post_train/sft/train_sft.py @@ -35,7 +35,6 @@ eval_interval=-1 steps=10 profiler=xplane """ -import inspect from typing import Any, Sequence from absl import app @@ -95,27 +94,26 @@ def create_train_step_fn(self): is_lora_enabled = self._lora_enabled wrt = nnx.LoRAParam if is_lora_enabled else nnx.Param - # Detect whether Tunix's train() expects (loss, aux, grad_norm) or just - # (loss, aux) by inspecting the source of PeftTrainer._train_step. - tunix_expects_grad_norm = False - try: - source = inspect.getsource(peft_trainer.PeftTrainer._train_step) # pylint: disable=protected-access - tunix_expects_grad_norm = "grad_norm" in source - except (TypeError, OSError): - pass - # Capture the graphdef once outside of JIT so that split/merge inside # jax.value_and_grad can use a stable (non-traced) structural descriptor. nnx.pop(self.model, nnx.Intermediate) graphdef, _, _ = nnx.split(self.model, wrt, ...) + _uses_gradient_accumulation = not ( + self.config.get_with_default("gradient_accumulation_steps", 1) == 1 and self.config.max_seq_token_per_tpu is None + ) def train_step( model: nnx.Module, optimizer: nnx.Optimizer, - grad_accumulator: Any, - inputs: Any, + grad_accumulator: Any = None, + inputs: Any = None, is_update_step: Any = True, ): + if inputs is None and grad_accumulator is not None: + # In Tunix versions where train_step only receives (model, optimizer, inputs) + inputs = grad_accumulator + grad_accumulator = getattr(self, "grad_accumulator", None) + inputs = gen_fn(inputs) # Split model into differentiable params and non-differentiable rest. @@ -160,12 +158,19 @@ def loss_wrapper(diff_params, rest, **inputs_kw): nnx.update(model, new_rest) # Handle gradient accumulation and conditional/direct optimizer update - if grad_accumulator is not None and hasattr(grad_accumulator, "add"): + if not _uses_gradient_accumulation: + if isinstance(grads, dict) and not isinstance(grads, nnx.State): + grads = nnx.State(grads) + optimizer.update(model, grads) + grad_norm = optax.global_norm(jax.tree_util.tree_map(lambda x: x.astype(jnp.float32), grads)) + else: grad_accumulator.add(grads) def apply_updates(model, optimizer, grad_accumulator): acc_grads = grad_accumulator.get() norm = optax.global_norm(jax.tree_util.tree_map(lambda x: x.astype(jnp.float32), acc_grads)) + if isinstance(acc_grads, dict) and not isinstance(acc_grads, nnx.State): + acc_grads = nnx.State(acc_grads) optimizer.update(model, acc_grads) grad_accumulator.reset() return norm @@ -181,14 +186,9 @@ def skip_updates(model, optimizer, grad_accumulator): optimizer, grad_accumulator, ) - else: - optimizer.update(model, grads) - grad_norm = optax.global_norm(jax.tree_util.tree_map(lambda x: x.astype(jnp.float32), grads)) aux_out = aux if has_aux else None - if tunix_expects_grad_norm: - return out_val, aux_out, grad_norm - return out_val, aux_out + return out_val, aux_out, grad_norm return train_step diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py index 4b313c178b..01e396d7a6 100644 --- a/src/maxtext/utils/lora_utils.py +++ b/src/maxtext/utils/lora_utils.py @@ -29,6 +29,23 @@ from orbax import checkpoint as ocp import qwix +try: + from qwix._src.utils import flax_util as _qwix_flax_util # pylint: disable=g-import-not-at-top + + _orig_qwix_find_param = _qwix_flax_util.find_param + + def _safe_qwix_find_param(x, ptq_array_type=None): + try: + return _orig_qwix_find_param(x, ptq_array_type) + except AttributeError as e: + if "shape" in str(e): + return None + raise + + _qwix_flax_util.find_param = _safe_qwix_find_param +except (ImportError, AttributeError): + pass + from maxtext.common import checkpointing from maxtext.configs import pyconfig from maxtext.utils import gcs_utils @@ -616,6 +633,9 @@ def apply_lora_to_model( decoder_positions=decoder_positions, rngs=model_rngs, ) + for _, val in nnx.graph.iter_graph(lora_model): + if hasattr(val, "__dict__") and "qwix_rngs" in val.__dict__: + del val.qwix_rngs if mesh is not None: with jax.set_mesh(mesh), nn_partitioning.axis_rules(mt_config.logical_axis_rules): @@ -640,6 +660,16 @@ def _safe_reshard(var, sharding_spec): val = var.get_value() if not isinstance(val, jax.Array): return var + if hasattr(sharding_spec, "spec") and len(sharding_spec.spec) != val.ndim: + spec_tuple = tuple(sharding_spec.spec) + if len(spec_tuple) > val.ndim: + if "local_layers" in spec_tuple and len(spec_tuple) - 1 == val.ndim: + spec_tuple = tuple(axis for axis in spec_tuple if axis != "local_layers") + else: + spec_tuple = spec_tuple[: val.ndim] + elif len(spec_tuple) < val.ndim: + spec_tuple = spec_tuple + (None,) * (val.ndim - len(spec_tuple)) + sharding_spec = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(*spec_tuple)) # make_array_from_callback natively constructs a globally sharded array # from the local host arrays, bypassing backend-specific device_put issues # on both Pathways and McJAX. diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index f162fc8275..8e97a09b1a 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -1950,8 +1950,12 @@ def _extract_primary_sharding(s): def _make_abstract_leaf(leaf_a, leaf_s): leaf_s = _extract_primary_sharding(leaf_s) - if hasattr(leaf_s, "spec") and len(leaf_a.shape) != len(leaf_s.spec): - leaf_s = jax.sharding.NamedSharding(leaf_s.mesh, jax.sharding.PartitionSpec(*leaf_s.spec[: len(leaf_a.shape)])) + if hasattr(leaf_s, "spec") and len(leaf_s.spec) > len(leaf_a.shape): + if "local_layers" in leaf_s.spec and len(leaf_s.spec) - 1 == len(leaf_a.shape): + spec_tuple = tuple(axis for axis in leaf_s.spec if axis != "local_layers") + else: + spec_tuple = leaf_s.spec[: len(leaf_a.shape)] + leaf_s = jax.sharding.NamedSharding(leaf_s.mesh, jax.sharding.PartitionSpec(*spec_tuple)) return jax.ShapeDtypeStruct(leaf_a.shape, leaf_a.dtype, sharding=leaf_s) if type(a_val) in (jax.Array, jax.ShapeDtypeStruct) or (hasattr(a_val, "shape") and not hasattr(a_val, "qvalue")): diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py index 3ebaa21610..1c4ff57b85 100644 --- a/src/maxtext/utils/sharding.py +++ b/src/maxtext/utils/sharding.py @@ -635,6 +635,18 @@ def add_data_to_sharding(mesh, path, aval, sharding): """ if not isinstance(sharding, jax.sharding.NamedSharding): raise AssertionError(f"Expected NamedSharding, found {sharding} of {type(sharding)=} at {jax.tree_util.keystr(path)}") + + pspec = sharding.spec + if len(pspec) != len(aval.shape): + if len(pspec) > len(aval.shape): + if "local_layers" in pspec and len(pspec) - 1 == len(aval.shape): + pspec = tuple(axis for axis in pspec if axis != "local_layers") + else: + pspec = pspec[: len(aval.shape)] + else: + pspec = tuple(pspec) + (None,) * (len(aval.shape) - len(pspec)) + sharding = jax.sharding.NamedSharding(sharding.mesh, jax.sharding.PartitionSpec(*pspec)) + try: sharded_shape = sharding.shard_shape(aval.shape) except Exception as e: diff --git a/tests/integration/lora_e2e_nnx_test.py b/tests/integration/lora_e2e_nnx_test.py index ee834fe225..5b43167510 100644 --- a/tests/integration/lora_e2e_nnx_test.py +++ b/tests/integration/lora_e2e_nnx_test.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Integration test for end-to-end Flax NNX LoRA checkpointing, resume, and adapter restoration across trainers.""" +"""Integration test for end-to-end Flax NNX LoRA checkpointing, resume, and adapter restoration.""" import os import shutil @@ -160,86 +160,6 @@ def _run_e2e_flow(self, model_name, use_sft, lora_weight_qtype=None, scan_layers state_step4 = train.train_loop(config_step4, recorder=None) self.assertEqual(int(state_step4.optimizer.step.get_value()), 2) - def _run_e2e_flow_sft(self, model_name, lora_weight_qtype=None, scan_layers=True): - """Executes a full 4-step E2E LoRA checkpoint/resume/restore flow for SFT (Tunix).""" - from maxtext.trainers.post_train.sft import train_sft # pylint: disable=import-outside-toplevel - - base_run_name = f"b_{model_name}_sft_run" - lora_run_name = f"w_{model_name}_sft_run" - - # Step 1: Generate base-only checkpoint (steps=2) - config_step1 = _tiny_lora_pyconfig( - run_name=base_run_name, - checkpoint_dir=self.test_dir, - model_name=model_name, - use_sft=True, - scan_layers=scan_layers, - steps=2, - checkpoint_period=2, - lora={"enable_lora": False}, - ) - trainer_step1, _ = train_sft.train(config_step1, goodput_recorder=None) - self.assertEqual(int(trainer_step1.train_steps), 2) - - base_ckpt_dir = os.path.join(self.test_dir, base_run_name, "checkpoints", "1") - self.assertTrue(os.path.exists(base_ckpt_dir), f"Base checkpoint path does not exist: {base_ckpt_dir}") - base_ckpt_path = os.path.join(base_ckpt_dir, "model_params") - - lora_config = {"enable_lora": True, "lora_rank": 4} - if lora_weight_qtype: - lora_config["lora_weight_qtype"] = lora_weight_qtype - lora_config["lora_tile_size"] = 4 - - # Step 2: Train with LoRA starting from base checkpoint (steps=4) - config_step2 = _tiny_lora_pyconfig( - run_name=lora_run_name, - checkpoint_dir=self.test_dir, - model_name=model_name, - use_sft=True, - scan_layers=scan_layers, - load_parameters_path=base_ckpt_path, - steps=4, - checkpoint_period=2, - lora=lora_config, - ) - trainer_step2, _ = train_sft.train(config_step2, goodput_recorder=None) - self.assertEqual(int(trainer_step2.train_steps), 4) - - lora_ckpt_dir = os.path.join(self.test_dir, lora_run_name, "checkpoints", "4") - self.assertTrue(os.path.exists(lora_ckpt_dir), f"Saved LoRA checkpoint path does not exist: {lora_ckpt_dir}") - lora_ckpt_path = os.path.join(lora_ckpt_dir, "model_params") - - # Step 3: Resume training under same run name (steps=6) - config_step3 = _tiny_lora_pyconfig( - run_name=lora_run_name, - checkpoint_dir=self.test_dir, - model_name=model_name, - use_sft=True, - scan_layers=scan_layers, - steps=6, - checkpoint_period=2, - lora=lora_config, - ) - trainer_step3, _ = train_sft.train(config_step3, goodput_recorder=None) - self.assertEqual(int(trainer_step3.train_steps), 6) - - # Step 4: Standalone restore of LoRA adapter onto base checkpoint (steps=2) - lora_restore_config = dict(lora_config) - lora_restore_config["lora_restore_path"] = lora_ckpt_path - config_step4 = _tiny_lora_pyconfig( - run_name=f"restore_{model_name}_sft_run", - checkpoint_dir=self.test_dir, - model_name=model_name, - use_sft=True, - scan_layers=scan_layers, - load_parameters_path=base_ckpt_path, - steps=2, - checkpoint_period=2, - lora=lora_restore_config, - ) - trainer_step4, _ = train_sft.train(config_step4, goodput_recorder=None) - self.assertEqual(int(trainer_step4.train_steps), 2) - # --- LoRA Unquantized Tests (Gemma4) --- def test_lora_e2e_gemma4_pretrain(self): self._run_e2e_flow("gemma4-26b", use_sft=False) @@ -247,10 +167,6 @@ def test_lora_e2e_gemma4_pretrain(self): def test_lora_e2e_gemma4_sft_native(self): self._run_e2e_flow("gemma4-26b", use_sft=True) - @pytest.mark.post_training - def test_lora_e2e_gemma4_sft(self): - self._run_e2e_flow_sft("gemma4-26b") - # --- QLoRA NF4 Tests (Gemma4, Qwen3, GPT-OSS) --- def test_qlora_e2e_gemma4_pretrain_nf4(self): self._run_e2e_flow("gemma4-26b", use_sft=False, lora_weight_qtype="nf4") @@ -258,20 +174,12 @@ def test_qlora_e2e_gemma4_pretrain_nf4(self): def test_qlora_e2e_gemma4_sft_native_nf4(self): self._run_e2e_flow("gemma4-26b", use_sft=True, lora_weight_qtype="nf4") - @pytest.mark.post_training - def test_qlora_e2e_gemma4_sft_nf4(self): - self._run_e2e_flow_sft("gemma4-26b", lora_weight_qtype="nf4") - def test_qlora_e2e_qwen3_pretrain_nf4(self): self._run_e2e_flow("qwen3-4b", use_sft=False, lora_weight_qtype="nf4") def test_qlora_e2e_qwen3_sft_native_nf4(self): self._run_e2e_flow("qwen3-4b", use_sft=True, lora_weight_qtype="nf4") - @pytest.mark.post_training - def test_qlora_e2e_qwen3_sft_nf4(self): - self._run_e2e_flow_sft("qwen3-4b", lora_weight_qtype="nf4") - def test_qlora_e2e_gptoss_unscanned_nf4(self): self._run_e2e_flow("gpt-oss-20b", use_sft=False, lora_weight_qtype="nf4", scan_layers=False) diff --git a/tests/post_training/integration/lora_e2e_nnx_test.py b/tests/post_training/integration/lora_e2e_nnx_test.py new file mode 100644 index 0000000000..63a664cd42 --- /dev/null +++ b/tests/post_training/integration/lora_e2e_nnx_test.py @@ -0,0 +1,175 @@ +# Copyright 2025-2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Integration test for end-to-end Flax NNX LoRA checkpointing, resume, and adapter restoration.""" + +import os +import shutil +import sys +import tempfile +import unittest + +from maxtext.configs import pyconfig +from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT +from tests.utils.test_helpers import get_test_config_path +import pytest + + +def _tiny_lora_pyconfig(run_name, checkpoint_dir, **overrides): + """Build a tiny pyconfig for E2E LoRA testing.""" + init_kwargs = { + "run_name": run_name, + "base_output_directory": checkpoint_dir, + "enable_checkpointing": True, + "dataset_type": "synthetic", + "model_name": "default", + "pure_nnx": True, + "per_device_batch_size": 1.0, + "base_emb_dim": 8, + "base_num_query_heads": 4, + "base_num_kv_heads": 4, + "base_mlp_dim": 32, + "base_num_decoder_layers": 2, + "head_dim": 128, + "max_target_length": 128, + "vocab_size": 256, + "steps": 10, + "async_checkpointing": False, + "checkpoint_period": 10, + "tokenizer_path": os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizers", "tokenizer.llama2"), + "enable_goodput_recording": False, + "enable_checkpoint_cloud_logger": False, + "monitor_goodput": False, + "override_model_config": True, + "use_tunix_gradient_accumulation": False, + "ici_fsdp_parallelism": 1, + "ici_tensor_parallelism": 1, + "ici_expert_parallelism": 1, + "ici_data_parallelism": -1, + "num_experts": 2, + "num_experts_per_tok": 1, + "shared_experts": 1, + "base_moe_mlp_dim": 32, + "attention": "dot_product", + } + init_kwargs.update(overrides) + return pyconfig.initialize([sys.argv[0], get_test_config_path()], **init_kwargs) + + +@pytest.mark.post_training +@pytest.mark.integration_test +class LoraE2ENnxIntegrationTest(unittest.TestCase): + """E2E integration test for NNX LoRA lifecycle. + + Covers base generation, LoRA train, resume, and standalone restore. + """ + + def setUp(self): + self.test_dir = tempfile.mkdtemp(prefix="lora_e2e_test_") + + def tearDown(self): + shutil.rmtree(self.test_dir, ignore_errors=True) + + def _run_e2e_flow_sft(self, model_name, lora_weight_qtype=None, scan_layers=True): + """Executes a full 4-step E2E LoRA checkpoint/resume/restore flow for SFT (Tunix).""" + from maxtext.trainers.post_train.sft import train_sft # pylint: disable=import-outside-toplevel + + base_run_name = f"b_{model_name}_sft_run" + lora_run_name = f"w_{model_name}_sft_run" + + # Step 1: Generate base-only checkpoint (steps=2) + config_step1 = _tiny_lora_pyconfig( + run_name=base_run_name, + checkpoint_dir=self.test_dir, + model_name=model_name, + use_sft=True, + scan_layers=scan_layers, + steps=2, + checkpoint_period=2, + lora={"enable_lora": False}, + ) + trainer_step1, _ = train_sft.train(config_step1, goodput_recorder=None) + self.assertEqual(int(trainer_step1.train_steps), 2) + + base_ckpt_dir = os.path.join(self.test_dir, base_run_name, "checkpoints", "2") + self.assertTrue(os.path.exists(base_ckpt_dir), f"Base checkpoint path does not exist: {base_ckpt_dir}") + base_ckpt_path = os.path.join(base_ckpt_dir, "model_params") + + lora_config = {"enable_lora": True, "lora_rank": 4} + if lora_weight_qtype: + lora_config["lora_weight_qtype"] = lora_weight_qtype + lora_config["lora_tile_size"] = 4 + + # Step 2: Train with LoRA starting from base checkpoint (steps=4) + config_step2 = _tiny_lora_pyconfig( + run_name=lora_run_name, + checkpoint_dir=self.test_dir, + model_name=model_name, + use_sft=True, + scan_layers=scan_layers, + load_parameters_path=base_ckpt_path, + steps=4, + checkpoint_period=2, + lora=lora_config, + ) + trainer_step2, _ = train_sft.train(config_step2, goodput_recorder=None) + self.assertEqual(int(trainer_step2.train_steps), 4) + + lora_ckpt_dir = os.path.join(self.test_dir, lora_run_name, "checkpoints", "4") + self.assertTrue(os.path.exists(lora_ckpt_dir), f"Saved LoRA checkpoint path does not exist: {lora_ckpt_dir}") + lora_ckpt_path = os.path.join(lora_ckpt_dir, "model_params") + + # Step 3: Resume training under same run name (steps=6) + config_step3 = _tiny_lora_pyconfig( + run_name=lora_run_name, + checkpoint_dir=self.test_dir, + model_name=model_name, + use_sft=True, + scan_layers=scan_layers, + steps=6, + checkpoint_period=2, + lora=lora_config, + ) + trainer_step3, _ = train_sft.train(config_step3, goodput_recorder=None) + self.assertEqual(int(trainer_step3.train_steps), 6) + + # Step 4: Standalone restore of LoRA adapter onto base checkpoint (steps=2) + lora_restore_config = dict(lora_config) + lora_restore_config["lora_restore_path"] = lora_ckpt_path + config_step4 = _tiny_lora_pyconfig( + run_name=f"restore_{model_name}_sft_run", + checkpoint_dir=self.test_dir, + model_name=model_name, + use_sft=True, + scan_layers=scan_layers, + load_parameters_path=base_ckpt_path, + steps=2, + checkpoint_period=2, + lora=lora_restore_config, + ) + trainer_step4, _ = train_sft.train(config_step4, goodput_recorder=None) + self.assertEqual(int(trainer_step4.train_steps), 2) + + def test_lora_e2e_gemma4_sft(self): + self._run_e2e_flow_sft("gemma4-26b") + + def test_qlora_e2e_gemma4_sft_nf4(self): + self._run_e2e_flow_sft("gemma4-26b", lora_weight_qtype="nf4") + + def test_qlora_e2e_qwen3_sft_nf4(self): + self._run_e2e_flow_sft("qwen3-4b", lora_weight_qtype="nf4") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/sharding_nnx_test.py b/tests/unit/sharding_nnx_test.py index c9f537b51b..e51676375d 100644 --- a/tests/unit/sharding_nnx_test.py +++ b/tests/unit/sharding_nnx_test.py @@ -284,11 +284,12 @@ def test_prevents_duplicate_physical_axes(self): ("embed", "fsdp"), ("mlp", "fsdp"), ) + v = nnx.Param( + jnp.zeros((3, 4)), + out_sharding=("embed", "mlp"), + eager_sharding=False, + ) with jax.set_mesh(self.mesh), nn_partitioning.axis_rules(rules): - v = nnx.Param( - jnp.zeros((3, 4)), - out_sharding=("embed", "mlp"), - ) out = self._run(self._build_state(w=v)) result_sharding = out["w"].get_value() self.assertIsInstance(result_sharding, NamedSharding) @@ -305,11 +306,12 @@ def test_fallback_to_next_physical_axis_when_duplicated(self): ("mlp", "fsdp"), ("mlp", "stage"), ) + v = nnx.Param( + jnp.zeros((3, 4)), + out_sharding=("embed", "mlp"), + eager_sharding=False, + ) with jax.set_mesh(self.mesh), nn_partitioning.axis_rules(rules): - v = nnx.Param( - jnp.zeros((3, 4)), - out_sharding=("embed", "mlp"), - ) out = self._run(self._build_state(w=v)) result_sharding = out["w"].get_value() self.assertIsInstance(result_sharding, NamedSharding) @@ -327,6 +329,7 @@ def test_resolves_when_context_rules_is_none(self): jnp.zeros((3,)), out_sharding=("embed",), sharding_rules=(("embed", "fsdp"),), + eager_sharding=False, ) out = self._run(self._build_state(w=v)) result_sharding = out["w"].get_value() @@ -354,13 +357,14 @@ def test_rules_merged_when_both_context_and_local_rules_present(self): # Local rules map 'embed' to 'stage'. Context rules map 'embed' to 'fsdp'. # Because local rules come first, 'embed' should resolve to 'stage'. context_rules = (("embed", "fsdp"),) + v = nnx.Param( + jnp.zeros((3,)), + out_sharding=("embed",), + sharding_rules=(("embed", "stage"),), + eager_sharding=False, + ) with jax.set_mesh(self.mesh), nn_partitioning.axis_rules(context_rules): - v = nnx.Param( - jnp.zeros((3,)), - out_sharding=("embed",), - sharding_rules=(("embed", "stage"),), - ) - out = self._run(self._build_state(w=v)) + out = self._run(self._build_state(w=v)) result_sharding = out["w"].get_value() self.assertEqual(result_sharding.spec, PartitionSpec("stage"))