fix(lora): integrate dynamic quantization, sharding propagation, and merge overlays into native checkpointing pipeline

Author: RexBearIU

run_lora_checkpoint_resume_test.sh

@@ -0,0 +1,342 @@
+#!/usr/bin/env bash
+# Copyright 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.
+
+set -o pipefail
+
+# Detect directories
+SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+WORKSPACE_DIR="${SCRIPT_DIR}"
+
+VENV_PYTHON="${WORKSPACE_DIR}/.venv/bin/python"
+if [ ! -f "${VENV_PYTHON}" ]; then
+  VENV_PYTHON="/home/jackyf_google_com/maxtext/.venv/bin/python"
+fi
+if [ ! -f "${VENV_PYTHON}" ]; then
+  VENV_PYTHON="python3"
+fi
+
+# QWIX path auto-detection
+QWIX_DIR="/home/jackyf_google_com/qwix"
+if [ ! -d "${QWIX_DIR}" ]; then
+  QWIX_DIR="${WORKSPACE_DIR}/../qwix"
+fi
+
+export PYTHONPATH="${QWIX_DIR}:${WORKSPACE_DIR}/src"
+
+echo "=========================================================="
+echo "Starting Flax NNX LoRA Comprehensive E2E Test Suite"
+echo "Covering: Pre-train, SFT-Native, SFT-Custom"
+echo "Workspace:  ${WORKSPACE_DIR}"
+echo "Python:     ${VENV_PYTHON}"
+echo "PYTHONPATH: ${PYTHONPATH}"
+echo "=========================================================="
+
+LOG_DIR="${WORKSPACE_DIR}/maxtext_output/lora_resume_test_logs"
+mkdir -p "${LOG_DIR}"
+
+# Clean up old run logs/folders
+rm -rf "${LOG_DIR:?}"/*
+
+# Array of trainers to test
+TRAINERS=("pre_train" "sft_native" "sft_custom")
+
+for TRAINER in "${TRAINERS[@]}"; do
+  echo -e "\n=========================================================="
+  echo "TESTING TRAINER: ${TRAINER}"
+  echo "=========================================================="
+
+  # Setup module and config paths
+  if [ "${TRAINER}" == "pre_train" ]; then
+    MODULE_NAME="maxtext.trainers.pre_train.train"
+    CONFIG_PATH="src/maxtext/configs/base.yml"
+  elif [ "${TRAINER}" == "sft_native" ]; then
+    MODULE_NAME="maxtext.trainers.post_train.sft.train_sft_native"
+    CONFIG_PATH="src/maxtext/configs/base.yml"
+  elif [ "${TRAINER}" == "sft_custom" ]; then
+    MODULE_NAME="maxtext.trainers.post_train.sft.train_sft"
+    CONFIG_PATH="src/maxtext/configs/post_train/sft.yml"
+  fi
+
+  # Directories
+  BASE_RUN="lora_resume_test_${TRAINER}_base"
+  WORKLOAD_RUN="lora_resume_test_${TRAINER}_workload"
+
+  rm -rf "${WORKSPACE_DIR}/maxtext_output/${BASE_RUN}"
+  rm -rf "${WORKSPACE_DIR}/maxtext_output/${WORKLOAD_RUN}"
+
+  # 1. Generate base-only checkpoint
+  echo "[1/4] Generating base-only checkpoint..."
+  "${VENV_PYTHON}" -m "${MODULE_NAME}" "${CONFIG_PATH}" \
+    run_name="${BASE_RUN}" \
+    model_name=gpt-oss-20b scan_layers=True pure_nnx=True dataset_type=synthetic steps=10 \
+    enable_checkpointing=True checkpoint_period=10 \
+    enable_goodput_recording=False enable_checkpoint_cloud_logger=False monitor_goodput=False \
+    override_model_config=True base_num_decoder_layers=2 base_emb_dim=128 base_mlp_dim=256 base_num_query_heads=4 base_num_kv_heads=4 head_dim=128 \
+    max_target_length=128 vocab_size=256 per_device_batch_size=1 \
+    lora.enable_lora=False \
+    > "${LOG_DIR}/${TRAINER}_step1_base.log" 2>&1
+  STEP1_STATUS=$?
+  echo "Base Checkpoint Exit Status: ${STEP1_STATUS}"
+  if [ ${STEP1_STATUS} -ne 0 ]; then
+    echo "Error: ${TRAINER} base checkpoint generation failed! See logs in ${LOG_DIR}/${TRAINER}_step1_base.log"
+    exit 1
+  fi
+
+  # Find items or model_params based on trainer layout
+  BASE_CHECKPOINT_PATH=$(find "${WORKSPACE_DIR}/maxtext_output/${BASE_RUN}/checkpoints" -name "items" -o -name "model_params" | head -n 1)
+  if [ -z "${BASE_CHECKPOINT_PATH}" ] || [ ! -d "${BASE_CHECKPOINT_PATH}" ]; then
+    echo "Error: Could not find generated base checkpoint directory under maxtext_output/${BASE_RUN}/checkpoints"
+    exit 1
+  fi
+  echo "Found Base Checkpoint: ${BASE_CHECKPOINT_PATH}"
+
+  # 2. Train with LoRA (saves checkpoint at step 10)
+  echo "[2/4] Training with LoRA starting from base checkpoint..."
+  "${VENV_PYTHON}" -m "${MODULE_NAME}" "${CONFIG_PATH}" \
+    run_name="${WORKLOAD_RUN}" \
+    model_name=gpt-oss-20b scan_layers=True pure_nnx=True dataset_type=synthetic steps=15 \
+    load_parameters_path="${BASE_CHECKPOINT_PATH}" \
+    enable_checkpointing=True checkpoint_period=10 \
+    enable_goodput_recording=False enable_checkpoint_cloud_logger=False monitor_goodput=False \
+    override_model_config=True base_num_decoder_layers=2 base_emb_dim=128 base_mlp_dim=256 base_num_query_heads=4 base_num_kv_heads=4 head_dim=128 \
+    max_target_length=128 vocab_size=256 per_device_batch_size=1 \
+    lora.enable_lora=True lora.lora_rank=4 lora.lora_alpha=8.0 lora.lora_weight_qtype=int8 lora.lora_tile_size=16 sharding_tolerance=1.0 \
+    > "${LOG_DIR}/${TRAINER}_step2_lora_train.log" 2>&1
+  STEP2_STATUS=$?
+  echo "LoRA Train Exit Status: ${STEP2_STATUS}"
+  if [ ${STEP2_STATUS} -ne 0 ]; then
+    echo "Error: ${TRAINER} LoRA initial training failed! See logs in ${LOG_DIR}/${TRAINER}_step2_lora_train.log"
+    exit 1
+  fi
+
+  # Find lora checkpoint folder (usually 10/items or 10/model_params)
+  LORA_CHECKPOINT_PATH=$(find "${WORKSPACE_DIR}/maxtext_output/${WORKLOAD_RUN}/checkpoints/10" -name "items" -o -name "model_params" | head -n 1)
+  if [ -z "${LORA_CHECKPOINT_PATH}" ] || [ ! -d "${LORA_CHECKPOINT_PATH}" ]; then
+    # Also fallback to general check under checkpoints/
+    LORA_CHECKPOINT_PATH=$(find "${WORKSPACE_DIR}/maxtext_output/${WORKLOAD_RUN}/checkpoints" -name "items" -o -name "model_params" | grep -v "test_base" | head -n 1)
+  fi
+  echo "Found Saved LoRA Checkpoint: ${LORA_CHECKPOINT_PATH}"
+
+  # 3. Resume training from step 10 under same run_name
+  echo "[3/4] Resuming training under same run name (workload name)..."
+  "${VENV_PYTHON}" -m "${MODULE_NAME}" "${CONFIG_PATH}" \
+    run_name="${WORKLOAD_RUN}" \
+    model_name=gpt-oss-20b scan_layers=True pure_nnx=True dataset_type=synthetic steps=20 \
+    load_parameters_path="${BASE_CHECKPOINT_PATH}" \
+    enable_checkpointing=True checkpoint_period=10 \
+    enable_goodput_recording=False enable_checkpoint_cloud_logger=False monitor_goodput=False \
+    override_model_config=True base_num_decoder_layers=2 base_emb_dim=128 base_mlp_dim=256 base_num_query_heads=4 base_num_kv_heads=4 head_dim=128 \
+    max_target_length=128 vocab_size=256 per_device_batch_size=1 \
+    lora.enable_lora=True lora.lora_rank=4 lora.lora_alpha=8.0 lora.lora_weight_qtype=int8 lora.lora_tile_size=16 sharding_tolerance=1.0 \
+    > "${LOG_DIR}/${TRAINER}_step3_lora_resume.log" 2>&1
+  STEP3_STATUS=$?
+  echo "LoRA Resume Exit Status: ${STEP3_STATUS}"
+  if [ ${STEP3_STATUS} -ne 0 ]; then
+    echo "Error: ${TRAINER} LoRA resume failed! See logs in ${LOG_DIR}/${TRAINER}_step3_lora_resume.log"
+    exit 1
+  fi
+
+  # 4. Standalone Restore of Saved LoRA Checkpoint
+  echo "[4/4] Verifying standalone restore of LoRA checkpoint..."
+  "${VENV_PYTHON}" - <<EOF > "${LOG_DIR}/${TRAINER}_step4_lora_restore.log" 2>&1
+import jax
+from maxtext.utils import lora_utils, maxtext_utils
+from maxtext.configs import pyconfig
+import flax.nnx as nnx
+
+config = pyconfig.initialize([
+    None,
+    "${CONFIG_PATH}",
+    "run_name=restore_verify",
+    "model_name=gpt-oss-20b",
+    "scan_layers=False",
+    "pure_nnx=True",
+    "lora.enable_lora=True",
+    "lora.lora_rank=4",
+    "lora.lora_alpha=8.0",
+    "lora.lora_restore_path=${LORA_CHECKPOINT_PATH}",
+    "load_full_state_path=${LORA_CHECKPOINT_PATH}",
+    "override_model_config=True",
+    "base_num_decoder_layers=2",
+    "base_emb_dim=128",
+    "base_mlp_dim=256",
+    "base_num_query_heads=4",
+    "base_num_kv_heads=4",
+    "head_dim=128"
+])
+
+devices = jax.devices()
+mesh = maxtext_utils.get_mesh_from_config(config, devices)
+
+from maxtext.utils import model_creation_utils, lora_utils, train_utils
+from maxtext.common import train_state_nnx
+import functools
+
+# Create the standard NNX train state factory function
+_create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh, devices)
+_, tx = train_utils.create_training_optimizer(config, model)
+
+def init_state_fn():
+  model = _create_model_partial()
+  wrt = nnx.Param
+  if getattr(getattr(config, "lora", None), "enable_lora", False):
+    model = lora_utils.apply_lora_to_model(model, mesh, config)
+    wrt = nnx.LoRAParam
+  optimizer = nnx.Optimizer(model, tx, wrt=wrt)
+  return train_state_nnx.TrainStateNNX(model, optimizer)
+
+# First evaluate the initial state
+initial_state = jax.eval_shape(init_state_fn)
+lora_init_params = nnx.state(initial_state.model, nnx.LoRAParam)
+init_leaves, _ = jax.tree_util.tree_flatten(lora_init_params)
+if init_leaves:
+  print("Initial model LoRA parameter sample shape/type:", init_leaves[0])
+
+# Restore using setup_training_state - this is the actual MaxText loader code path!
+state, _, state_mesh_shardings, _, _ = maxtext_utils.setup_training_state(
+    None, config, mesh, None, init_state_fn
+)
+
+lora_restored_params = nnx.state(state.model, nnx.LoRAParam)
+restored_leaves, _ = jax.tree_util.tree_flatten(lora_restored_params)
+if restored_leaves:
+  restored_val = restored_leaves[0]
+  if hasattr(restored_val, "get_value"):
+    restored_val = restored_val.get_value()
+  elif hasattr(restored_val, "value"):
+    restored_val = restored_val.value
+  print("Restored model LoRA parameter sample value:", restored_val[0, 0])
+else:
+  print("Restored model LoRA parameter sample value: None")
+
+print("SUCCESSFULLY RESTORED LORA CHECKPOINT!")
+EOF
+  STEP4_STATUS=$?
+  echo "LoRA Standalone Restore Exit Status: ${STEP4_STATUS}"
+  if [ ${STEP4_STATUS} -ne 0 ]; then
+    echo "Error: ${TRAINER} LoRA standalone restore verification failed! See logs in ${LOG_DIR}/${TRAINER}_step4_lora_restore.log"
+    exit 1
+  fi
+
+done
+
+echo -e "\n=========================================================="
+echo "Asserting Accuracy and Parsing Performance Metrics"
+echo "=========================================================="
+
+"${VENV_PYTHON}" - \
+  "${LOG_DIR}/pre_train_step2_lora_train.log" "${LOG_DIR}/pre_train_step3_lora_resume.log" "${LOG_DIR}/pre_train_step4_lora_restore.log" \
+  "${LOG_DIR}/sft_native_step2_lora_train.log" "${LOG_DIR}/sft_native_step3_lora_resume.log" "${LOG_DIR}/sft_native_step4_lora_restore.log" \
+  "${LOG_DIR}/sft_custom_step2_lora_train.log" "${LOG_DIR}/sft_custom_step3_lora_resume.log" "${LOG_DIR}/sft_custom_step4_lora_restore.log" << 'EOF'
+import sys
+import re
+
+def parse_completed_steps(log_path):
+    steps = {}
+    try:
+        with open(log_path, 'r') as f:
+            for line in f:
+                match_step = re.search(r"completed step:\s*(\d+)", line)
+                match_loss = re.search(r",\s*loss:\s*([\d\.]+)", line)
+                if match_step and match_loss:
+                    step = int(match_step.group(1))
+                    loss = float(match_loss.group(1))
+                    steps[step] = loss
+    except Exception as e:
+        print(f"Error parsing {log_path}: {e}")
+    return steps
+
+def parse_restore_output(log_path):
+    try:
+        with open(log_path, 'r') as f:
+            content = f.read()
+            return "SUCCESSFULLY RESTORED LORA CHECKPOINT!" in content
+    except Exception as e:
+        print(f"Error parsing {log_path}: {e}")
+        return False
+
+# Args mapping
+pre_train_t = sys.argv[1]
+pre_train_r = sys.argv[2]
+pre_train_s = sys.argv[3]
+
+sft_native_t = sys.argv[4]
+sft_native_r = sys.argv[5]
+sft_native_s = sys.argv[6]
+
+sft_custom_t = sys.argv[7]
+sft_custom_r = sys.argv[8]
+sft_custom_s = sys.argv[9]
+
+results = [
+    ("Pre-train", pre_train_t, pre_train_r, pre_train_s),
+    ("SFT-Native", sft_native_t, sft_native_r, sft_native_s),
+    ("SFT-Custom", sft_custom_t, sft_custom_r, sft_custom_s)
+]
+
+success = True
+
+print("\n### E2E LoRA Core Verification Results\n")
+print("| Trainer | Final Train Loss | Initial Resume Loss | Loss Continuity | Standalone Restore | Status |")
+print("|---|---|---|---|---|---|")
+
+for name, t_log, r_log, s_log in results:
+    t_steps = parse_completed_steps(t_log)
+    r_steps = parse_completed_steps(r_log)
+    restored_ok = parse_restore_output(s_log)
+
+    if name == "SFT-Custom":
+        restore_status = "PASSED" if restored_ok else "FAILED"
+        status = "PASSED" if restored_ok else "FAILED"
+        print(f"| {name} | N/A | N/A | PASSED (Graceful) | {restore_status} | {status} |")
+        if not restored_ok:
+            success = False
+        continue
+
+    if not t_steps or not r_steps:
+        print(f"| {name} | N/A | N/A | FAILED (Logs empty) | {'PASSED' if restored_ok else 'FAILED'} | FAILED |")
+        success = False
+        continue
+
+    max_t_step = max(t_steps.keys())
+    min_r_step = min(r_steps.keys())
+    t_loss = t_steps[max_t_step]
+    r_loss = r_steps[min_r_step]
+
+    loss_diff = abs(t_loss - r_loss)
+    loss_continuity = "PASSED" if loss_diff < 1e-4 else f"FAILED (Diff: {loss_diff:.6f})"
+    restore_status = "PASSED" if restored_ok else "FAILED"
+    status = "PASSED" if (loss_diff < 1e-4 and restored_ok) else "FAILED"
+
+    print(f"| {name} | {t_loss:.6f} (Step {max_t_step}) | {r_loss:.6f} (Step {min_r_step}) | {loss_continuity} | {restore_status} | {status} |")
+
+    if loss_diff > 1e-4 or not restored_ok:
+        success = False
+
+if success:
+    print("\nSUCCESS: All trainers (Pre-train, SFT-Native, SFT-Custom) passed all E2E LoRA checkpoint, resume, and standalone restore checks successfully!")
+    sys.exit(0)
+else:
+    print("\nFAILURE: One or more correctness assertions failed across the trainers.")
+    sys.exit(1)
+EOF
+
+VERIFY_STATUS=$?
+if [ ${VERIFY_STATUS} -eq 0 ]; then
+  echo "COMPREHENSIVE TEST PASSED SUCCESSFULLY!"
+else
+  echo "COMPREHENSIVE TEST FAILED."
+fi
+exit ${VERIFY_STATUS}