From a432433a387e12d9c7a8e0ebb5a803e205514a35 Mon Sep 17 00:00:00 2001 From: Jacky Fang Date: Thu, 23 Jul 2026 09:05:11 +0000 Subject: [PATCH] feat(nnx): Add Flax NNX PEFT/LoRA support and checkpoint resume - Resolve pjit sharding metadata mismatches by re-evaluating state_mesh_shardings after parameter restoration and LoRA/QLoRA module injection. - Update l2norm_pytree to skip non-inexact leaves (e.g. packed void arrays or WithAux wrappers). - Implement normalized key translation (0 <-> layers_0) in checkpointing and model creation utilities for seamless parameter merging across NNX and Linen models. - Add mesh axis shape sanitization for non-divisible dimensions. - Include unit tests for l2norm_pytree void quantization, key translation, and upfront LoRA state restoration. --- run_lora_checkpoint_resume_test.sh | 321 ++++++ src/maxtext/common/checkpointing.py | 735 ++++++++++---- src/maxtext/common/train_state_nnx.py | 180 +++- .../configs/post_train/lora_module_path.yml | 1 + src/maxtext/trainers/pre_train/train.py | 326 ++++-- src/maxtext/utils/gradient_accumulation.py | 192 ++-- src/maxtext/utils/lora_utils.py | 462 ++++++--- src/maxtext/utils/max_utils.py | 367 ++++--- src/maxtext/utils/maxtext_utils.py | 957 ++++++++++-------- src/maxtext/utils/sharding.py | 264 +++-- src/maxtext/utils/train_utils.py | 34 +- tests/integration/peft_integration_test.py | 654 ++++++++++++ .../integration/setup_train_loop_nnx_test.py | 15 +- tests/post_training/unit/lora_utils_test.py | 245 ++++- tests/unit/checkpointing_test.py | 90 +- tests/unit/maxtext_utils_test.py | 11 +- tests/unit/model_creation_utils_test.py | 326 ++++-- 17 files changed, 3814 insertions(+), 1366 deletions(-) create mode 100755 run_lora_checkpoint_resume_test.sh create mode 100644 tests/integration/peft_integration_test.py diff --git a/run_lora_checkpoint_resume_test.sh b/run_lora_checkpoint_resume_test.sh new file mode 100755 index 0000000000..46c7cb918c --- /dev/null +++ b/run_lora_checkpoint_resume_test.sh @@ -0,0 +1,321 @@ +#!/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="/home/jackyf_google_com/maxtext/.venv/bin/python" +if [ ! -f "${VENV_PYTHON}" ]; then + VENV_PYTHON="${WORKSPACE_DIR}/.venv/bin/python" +fi +if [ ! -f "${VENV_PYTHON}" ]; then + VENV_PYTHON="python3" +fi + +export PYTHONPATH="${WORKSPACE_DIR}/src" + +# Parse CLI arguments and options +SPECIFIED_MODEL="" +SPECIFIED_TRAINER="" +SPECIFIED_STEP="all" +SCAN_LAYERS_VAL="True" +LORA_QTYPE="nf4" + +while [[ $# -gt 0 ]]; do + case "$1" in + -m|--model) + SPECIFIED_MODEL="$2" + shift 2 + ;; + -t|--trainer) + SPECIFIED_TRAINER="$2" + shift 2 + ;; + -s|--step) + SPECIFIED_STEP="$2" + shift 2 + ;; + --scan|--scan-layers) + SCAN_LAYERS_VAL="$2" + shift 2 + ;; + -q|--qtype) + LORA_QTYPE="$2" + shift 2 + ;; + -h|--help) + echo "Usage: $0 [-m|--model MODEL_NAME] [-t|--trainer TRAINER_TYPE] [-s|--step STEP_NUM] [--scan True|False] [-q|--qtype nf4|int8]" + exit 0 + ;; + *) + if [ -z "${SPECIFIED_MODEL}" ]; then + SPECIFIED_MODEL="$1" + elif [ -z "${SPECIFIED_TRAINER}" ]; then + SPECIFIED_TRAINER="$1" + fi + shift + ;; + esac +done + +# Fallback to environment variables if set +SPECIFIED_MODEL="${SPECIFIED_MODEL:-${TEST_MODEL_NAME}}" +SPECIFIED_TRAINER="${SPECIFIED_TRAINER:-${TEST_TRAINER}}" + +if [ -n "${SPECIFIED_MODEL}" ]; then + MODELS=("${SPECIFIED_MODEL}") +else + MODELS=("qwen3-4b" "gemma3-4b" "llama3.1-8b" "gpt-oss-20b" "gemma4-26b") +fi + +if [ -n "${SPECIFIED_TRAINER}" ]; then + TRAINERS=("${SPECIFIED_TRAINER}") +else + TRAINERS=("pre_train" "sft_native" "sft_custom") +fi + +echo "==========================================================" +echo "Starting Flax NNX LoRA Comprehensive E2E Test Suite" +echo "Models under test: ${MODELS[*]}" +echo "Trainers under test:${TRAINERS[*]}" +echo "Scan Layers: ${SCAN_LAYERS_VAL}" +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}" + +for TEST_MODEL_NAME in "${MODELS[@]}"; do +for TRAINER in "${TRAINERS[@]}"; do + echo -e "\n==========================================================" + echo "TESTING MODEL: ${TEST_MODEL_NAME} | 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_${TEST_MODEL_NAME}_${TRAINER}_base" + WORKLOAD_RUN="lora_resume_test_${TEST_MODEL_NAME}_${TRAINER}_workload" + + if [ "${SPECIFIED_STEP}" = "all" ] || [ "${SPECIFIED_STEP}" = "1" ]; then + rm -rf "${WORKSPACE_DIR}/maxtext_output/${BASE_RUN}" + rm -rf "${WORKSPACE_DIR}/maxtext_output/${WORKLOAD_RUN}" + elif [ "${SPECIFIED_STEP}" = "2" ]; then + rm -rf "${WORKSPACE_DIR}/maxtext_output/${WORKLOAD_RUN}" + fi + + LOG_PREFIX="${LOG_DIR}/${TEST_MODEL_NAME}_${TRAINER}" + + # 1. Generate base-only checkpoint + if [ "${SPECIFIED_STEP}" = "all" ] || [ "${SPECIFIED_STEP}" = "1" ]; then + echo "[1/4] Generating base-only checkpoint..." + "${VENV_PYTHON}" -m "${MODULE_NAME}" "${CONFIG_PATH}" \ + run_name="${BASE_RUN}" \ + model_name="${TEST_MODEL_NAME}" scan_layers=${SCAN_LAYERS_VAL} 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 log_period=1 \ + 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=8 \ + lora.enable_lora=False \ + > "${LOG_PREFIX}_step1_base.log" 2>&1 + STEP1_STATUS=$? + echo "Base Checkpoint Exit Status: ${STEP1_STATUS}" + if [ ${STEP1_STATUS} -ne 0 ]; then + echo "Error: ${TEST_MODEL_NAME} ${TRAINER} base checkpoint generation failed! See logs in ${LOG_PREFIX}_step1_base.log" + exit 1 + fi + 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) + if [ "${SPECIFIED_STEP}" = "all" ] || [ "${SPECIFIED_STEP}" = "2" ]; then + sleep 2 + echo "[2/4] Training with LoRA starting from base checkpoint..." + "${VENV_PYTHON}" -m "${MODULE_NAME}" "${CONFIG_PATH}" \ + run_name="${WORKLOAD_RUN}" \ + model_name="${TEST_MODEL_NAME}" scan_layers=${SCAN_LAYERS_VAL} 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 log_period=1 \ + 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=8 \ + lora.enable_lora=True lora.lora_rank=4 lora.lora_alpha=8.0 lora.lora_weight_qtype=${LORA_QTYPE} lora.lora_tile_size=16 sharding_tolerance=1.0 \ + > "${LOG_PREFIX}_step2_lora_train.log" 2>&1 + STEP2_STATUS=$? + echo "LoRA Train Exit Status: ${STEP2_STATUS}" + if [ ${STEP2_STATUS} -ne 0 ]; then + echo "Error: ${TEST_MODEL_NAME} ${TRAINER} LoRA initial training failed! See logs in ${LOG_PREFIX}_step2_lora_train.log" + exit 1 + fi + fi + + # Find lora checkpoint folder (any items or model_params folder under checkpoints) + LORA_CHECKPOINT_PATH=$(find "${WORKSPACE_DIR}/maxtext_output/${WORKLOAD_RUN}/checkpoints" -type d \( -name "items" -o -name "model_params" \) | grep -v "/0/" | head -n 1) + if [ -z "${LORA_CHECKPOINT_PATH}" ]; then + LORA_CHECKPOINT_PATH=$(find "${WORKSPACE_DIR}/maxtext_output/${WORKLOAD_RUN}/checkpoints" -type d \( -name "items" -o -name "model_params" \) | head -n 1) + fi + echo "Found Saved LoRA Checkpoint: ${LORA_CHECKPOINT_PATH}" + + # 3. Resume training from step 10 under same run_name + if [ "${SPECIFIED_STEP}" = "all" ] || [ "${SPECIFIED_STEP}" = "3" ]; then + sleep 2 + echo "[3/4] Resuming training under same run name (workload name)..." + "${VENV_PYTHON}" -m "${MODULE_NAME}" "${CONFIG_PATH}" \ + run_name="${WORKLOAD_RUN}" \ + model_name="${TEST_MODEL_NAME}" scan_layers=${SCAN_LAYERS_VAL} pure_nnx=True dataset_type=synthetic steps=20 \ + enable_checkpointing=True checkpoint_period=10 \ + enable_goodput_recording=False enable_checkpoint_cloud_logger=False monitor_goodput=False log_period=1 \ + 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=8 \ + lora.enable_lora=True lora.lora_rank=4 lora.lora_alpha=8.0 lora.lora_weight_qtype=${LORA_QTYPE} lora.lora_tile_size=16 sharding_tolerance=1.0 \ + > "${LOG_PREFIX}_step3_lora_resume.log" 2>&1 + STEP3_STATUS=$? + echo "LoRA Resume Exit Status: ${STEP3_STATUS}" + if [ ${STEP3_STATUS} -ne 0 ]; then + echo "Error: ${TEST_MODEL_NAME} ${TRAINER} LoRA resume training failed! See logs in ${LOG_PREFIX}_step3_lora_resume.log" + exit 1 + fi + fi + + # 4. Verify standalone restore of saved LoRA checkpoint + if [ "${SPECIFIED_STEP}" = "all" ] || [ "${SPECIFIED_STEP}" = "4" ]; then + sleep 2 + echo "[4/4] Verifying standalone restore of LoRA checkpoint..." + "${VENV_PYTHON}" -m "${MODULE_NAME}" "${CONFIG_PATH}" \ + run_name="lora_restore_verify_${TEST_MODEL_NAME}_${TRAINER}_$(date +%s)" \ + model_name="${TEST_MODEL_NAME}" scan_layers=${SCAN_LAYERS_VAL} pure_nnx=True dataset_type=synthetic steps=11 \ + load_parameters_path="${BASE_CHECKPOINT_PATH}" \ + lora.lora_restore_path="${LORA_CHECKPOINT_PATH}" \ + enable_checkpointing=True checkpoint_period=1000 \ + enable_goodput_recording=False enable_checkpoint_cloud_logger=False monitor_goodput=False log_period=1 \ + 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=8 \ + lora.enable_lora=True lora.lora_rank=4 lora.lora_alpha=8.0 lora.lora_weight_qtype=${LORA_QTYPE} lora.lora_tile_size=16 sharding_tolerance=1.0 \ + > "${LOG_PREFIX}_step4_lora_restore.log" 2>&1 + STEP4_STATUS=$? + echo "LoRA Standalone Restore Exit Status: ${STEP4_STATUS}" + if [ ${STEP4_STATUS} -ne 0 ]; then + echo "Error: ${TEST_MODEL_NAME} ${TRAINER} LoRA standalone restore failed! See logs in ${LOG_PREFIX}_step4_lora_restore.log" + exit 1 + fi + fi + +done +done + +echo "" +echo "==========================================================" +echo "Asserting Accuracy and Parsing Performance Metrics" +echo "==========================================================" + +"${VENV_PYTHON}" -c " +import glob +import re +import os + +log_dir = '${LOG_DIR}' +models = [$(printf '"%s", ' "${MODELS[@]}")] +trainers = [$(printf '"%s", ' "${TRAINERS[@]}")] + +# Find all log files +log_files = glob.glob(os.path.join(log_dir, '*.log')) + +results = [] +all_passed = True + +def parse_loss(filepath, target_step): + if not os.path.exists(filepath): + return None, f'File not found: {os.path.basename(filepath)}' + with open(filepath, 'r') as f: + content = f.read() + pattern = rf'completed step: {target_step},.*loss: ([\d\.]+)' + match = re.search(pattern, content) + if match: + return float(match.group(1)), None + return None, f'Step {target_step} not found in {os.path.basename(filepath)}' + +print('\n### E2E LoRA Core Verification Results\n') +print('| Model | Trainer | Final Train Loss | Initial Resume Loss | Loss Continuity | Standalone Restore | Status |') +print('|---|---|---|---|---|---|---|') + +for model in models: + for trainer in trainers: + step2_log = os.path.join(log_dir, f'{model}_{trainer}_step2_lora_train.log') + step3_log = os.path.join(log_dir, f'{model}_{trainer}_step3_lora_resume.log') + step4_log = os.path.join(log_dir, f'{model}_{trainer}_step4_lora_restore.log') + + # SFT trainers save/resume at step 10/11 vs pre_train step 10/11 + train_step = 15 if trainer == 'sft_custom' else 14 + resume_step = 16 if trainer == 'sft_custom' else 15 + + train_loss, err2 = parse_loss(step2_log, train_step) + resume_loss, err3 = parse_loss(step3_log, resume_step) + + restore_passed = os.path.exists(step4_log) + if restore_passed: + with open(step4_log, 'r') as f: + restore_content = f.read() + restore_passed = 'completed step: 10' in restore_content or 'completed step: 0' in restore_content + + if train_loss is not None and resume_loss is not None: + diff = abs(train_loss - resume_loss) + rel_diff = diff / max(abs(train_loss), 1e-5) + loss_passed = diff < 0.05 or rel_diff < 0.05 + continuity_str = f'PASSED' if loss_passed else f'FAILED (Diff: {diff:.6f}, RelDiff: {rel_diff:.4f})' + else: + loss_passed = False + continuity_str = 'FAILED (Logs empty)' + + restore_str = 'PASSED' if restore_passed else 'FAILED' + row_passed = loss_passed and restore_passed + if not row_passed: + all_passed = False + + status_str = 'PASSED' if row_passed else 'FAILED' + train_str = f'{train_loss:.6f} (Step {train_step})' if train_loss else 'N/A' + resume_str = f'{resume_loss:.6f} (Step {resume_step})' if resume_loss else 'N/A' + + print(f'| {model} | {trainer} | {train_str} | {resume_str} | {continuity_str} | {restore_str} | {status_str} |') + +if not all_passed: + print('\nFAILURE: One or more correctness assertions failed across the models/trainers.') + exit(1) + +print('\nSUCCESS: All Flax NNX LoRA E2E correctness assertions passed successfully!') +" + +echo "COMPREHENSIVE TEST COMPLETE." \ No newline at end of file diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index cf51083597..80be9b82a5 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -11,8 +11,6 @@ # 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. - - """Create an Orbax CheckpointManager with specified (Async or not) Checkpointer.""" import datetime @@ -34,7 +32,14 @@ from maxtext.utils import elastic_utils from maxtext.utils import exceptions from maxtext.utils import gcs_utils +from maxtext.utils import lora_utils from maxtext.utils import max_logging +import jax.numpy as jnp +import numpy as np + +from qwix import QArray +from qwix._src.providers.ptq import WithAux + from maxtext.utils.globals import DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE import orbax.checkpoint as ocp from orbax.checkpoint import v1 as ocp_v1 @@ -42,7 +47,6 @@ from orbax.checkpoint._src.checkpoint_managers import preservation_policy as preservation_policy_lib from orbax.checkpoint._src.checkpoint_managers import save_decision_policy as save_decision_policy_lib - CheckpointManagerOptions = ocp.CheckpointManagerOptions Composite = ocp.args.Composite PyTreeCheckpointHandler = ocp.PyTreeCheckpointHandler @@ -53,28 +57,305 @@ create_orbax_emergency_replicator_checkpoint_manager = emergency_checkpointing.create_replicator_checkpoint_manager # Union of CheckpointManager / the emergency factories return; used in type hints. +import flax + CheckpointManager = ocp.CheckpointManager | EmergencyCheckpointManager | EmergencyReplicatorCheckpointManager -def _weight_mismatches(want, have, path=()): +def _tree_to_dict(tree): + """Recursively converts NNX State or PyTree to pure python dict.""" + if hasattr(tree, "to_pure_dict"): + return _tree_to_dict(tree.to_pure_dict()) + if isinstance(tree, dict): + res = {k: _tree_to_dict(v) for k, v in tree.items()} + if len(res) == 1 and ("value" in res or "raw_value" in res): + return res.get("value", res.get("raw_value")) + return res + if hasattr(tree, "get_value"): + return tree.get_value() + if hasattr(tree, "value"): + return tree.value + return tree + + +def _norm_path_key(path): + """Normalizes a PyTree path key or string path to a canonical string representation.""" + if isinstance(path, str): + parts = path.split("/") + elif isinstance(path, (tuple, list)): + parts = [str(p) for p in path] + else: + parts = [str(path)] + if len(parts) > 1 and parts[-1] in ("value", "raw_value"): + parts = parts[:-1] + norm_parts = [] + for p in parts: + if p in ("scanned_blocks", "layers", "layers_remainder", "model", "params"): + continue + if p.isdigit(): + norm_parts.append(f"layers_{p}") + else: + norm_parts.append(p) + return "/".join(norm_parts) + + +def _is_dict_leaf(x): + if not isinstance(x, (dict, nnx.State)): + return True + if ("qvalue" in x and "scale" in x) or ("array" in x and "how" in x): + return True + if "array" in x and isinstance( + x["array"], + (dict, nnx.State)) and ("qvalue" in x["array"] and "scale" in x["array"]): + return True + return False + + +def _flatten_and_norm_dict(d): + """Flattens dict d with flax.traverse_util.flatten_dict and normalizes path keys with _norm_path_key.""" + if d is None: + return {} + d = _tree_to_dict(d) + flat = flax.traverse_util.flatten_dict( + d, + is_leaf=lambda *path: _is_dict_leaf(path[-1]), + ) + return {_norm_path_key(k): v for k, v in flat.items()} + + +def _weight_mismatches(want, have): """Returns `(path, problem)` for each weight in `want` that `have` didn't restore faithfully. A weight is wrong if the checkpoint didn't carry it -- absent, or left by Orbax as an unmaterialized ShapeDtypeStruct -- or carried it at a different shape. Only the shape can disagree: Orbax casts a restored array to the target's dtype. + + For PEFT/LoRA and QLoRA compatibility: + - Runtime-injected LoRA parameters (`kernel_lora_a`, `kernel_lora_b`) and RNG states are skipped. + - Dynamically quantized QLoRA weights expecting `qvalue` / `scale` leaves are matched against + unquantized base parameter keys (e.g. `kernel`). """ - if isinstance(want, dict): - out = [] - for k, v in want.items(): - out.extend(_weight_mismatches(v, have.get(k) if isinstance(have, dict) else None, path + (k,))) - return out - name = "/".join(str(p) for p in path) - if have is None or isinstance(have, jax.ShapeDtypeStruct): - return [(name, f"missing (model expects {getattr(want, 'shape', '?')} {getattr(want, 'dtype', '?')})")] - want_shape, got_shape = getattr(want, "shape", None), getattr(have, "shape", None) - if want_shape is not None and got_shape is not None and tuple(want_shape) != tuple(got_shape): - return [(name, f"shape {tuple(got_shape)} but the model expects {tuple(want_shape)}")] - return [] + if want is None: + return [] + flat_want = flax.traverse_util.flatten_dict(_tree_to_dict(want)) + norm_have = _flatten_and_norm_dict(have) + + problems = [] + for path, target_val in flat_want.items(): + while isinstance(target_val, nnx.Variable): + target_val = target_val.get_value() if hasattr( + target_val, "get_value") else target_val.value + + name = "/".join(str(p) for p in path) + path_parts = [str(p) for p in path] + norm_p = _norm_path_key(path) + restored_val = norm_have.get(norm_p) + + if "lora_a" in name or "lora_b" in name or "rngs" in path_parts or "rng" in path_parts: + if restored_val is None or isinstance(restored_val, jax.ShapeDtypeStruct): + continue + + # If quantized QArray leaf is missing, check if unquantized base parameter exists in restored checkpoint + if (restored_val is None or + isinstance(restored_val, jax.ShapeDtypeStruct)) and any( + k in name for k in ("qvalue", "scale", "qarray", "zero_point")): + base_parts = [ + p for p in path_parts if p not in ("array", "qvalue", "scale", + "qarray", "zero_point", "bits") + ] + base_norm_p = _norm_path_key(base_parts) + if base_norm_p in norm_have and not isinstance(norm_have[base_norm_p], + jax.ShapeDtypeStruct): + continue + + if restored_val is None or isinstance(restored_val, jax.ShapeDtypeStruct): + target_shape = getattr(target_val, "shape", "?") + target_dtype = getattr(target_val, "dtype", "?") + problems.append( + (name, f"missing (model expects {target_shape} {target_dtype})")) + else: + want_shape, got_shape = getattr(target_val, "shape", + None), getattr(restored_val, "shape", + None) + if want_shape is not None and got_shape is not None and tuple( + want_shape) != tuple(got_shape): + problems.append(( + name, + f"shape {tuple(got_shape)} but the model expects {tuple(want_shape)}" + )) + return problems + + +def _set_nested_leaf(target, path, leaf_val): + """Sets a nested leaf in target given a path.""" + curr = target + for key_obj in path[:-1]: + k = getattr(key_obj, "key", + str(key_obj)) if not isinstance(key_obj, + (str, int)) else key_obj + if isinstance(curr, (dict, nnx.State)) or hasattr(curr, "__getitem__"): + curr = curr[k] + elif hasattr(curr, str(k)): + curr = getattr(curr, str(k)) + else: + return + last_k = getattr(path[-1], "key", str( + path[-1])) if not isinstance(path[-1], (str, int)) else path[-1] + if isinstance(curr, (dict, nnx.State)) or hasattr(curr, "__setitem__"): + curr[last_k] = leaf_val + elif hasattr(curr, str(last_k)): + setattr(curr, str(last_k), leaf_val) + + +def _rebuild_qwix_types(val): + """Recursively reconstructs QArray or WithAux objects from dict/State representations restored from Orbax.""" + if hasattr(val, "to_pure_dict"): + val = val.to_pure_dict() + if isinstance(val, (dict, nnx.State)): + d = {k: _rebuild_qwix_types(v) for k, v in val.items()} + if "qvalue" in d and "scale" in d: + qval = d["qvalue"] + scale = d["scale"] + qval = qval.get_value() if hasattr(qval, "get_value") else getattr( + qval, "value", qval) + scale = scale.get_value() if hasattr(scale, "get_value") else getattr( + scale, "value", scale) + if hasattr(qval, "dtype") and jnp.issubdtype(qval.dtype, jnp.floating): + qval = qval.astype(jnp.int8) + zp = d.get("zero_point") + if zp is not None: + zp = zp.get_value() if hasattr(zp, "get_value") else getattr( + zp, "value", zp) + qtype = d.get("qtype", "nf4") + if isinstance(qtype, (jax.Array, np.ndarray)) or hasattr(qtype, "item"): + qtype = str(qtype.item()) if hasattr(qtype, "item") else str(qtype) + qarr = QArray(qvalue=qval, scale=scale, zero_point=zp, qtype=qtype) + return WithAux(array=qarr, how=d.get("how", "ptq")) + if "array" in d: + arr = d["array"] + arr = arr.get_value() if hasattr(arr, "get_value") else getattr( + arr, "value", arr) + how = d.get("how", "ptq") + how = how.get_value() if hasattr(how, "get_value") else getattr( + how, "value", how) + return WithAux(array=arr, how=how) + return d + if hasattr(val, "get_value"): + return _rebuild_qwix_types(val.get_value()) + if hasattr(val, "value"): + return _rebuild_qwix_types(val.value) + return val + + +def _update_leaf_var(var, leaf_val, target_root, path): + """Updates an NNX Variable or target leaf with the given leaf value.""" + leaf_val = _rebuild_qwix_types(leaf_val) + target_sharding = getattr(var, "sharding", None) + if target_sharding is None and hasattr(var, "get_value"): + target_sharding = getattr(var.get_value(), "sharding", None) + + if (target_sharding is not None and + isinstance(target_sharding, jax.sharding.Sharding) and + not isinstance(leaf_val, jax.ShapeDtypeStruct)): + leaf_val = jax.device_put(leaf_val, target_sharding) + + if isinstance(leaf_val, (jax.Array, np.ndarray)): + leaf_val = jnp.copy(leaf_val) + + if hasattr(var, "set_value"): + var.set_value(leaf_val) + elif hasattr(var, "value"): + var.value = leaf_val + else: + _set_nested_leaf(target_root, path, leaf_val) + + +def _norm_nnx_path_key(path_tuple, root_state=None): + """Normalize NNX path key tuple into a canonical slash-delimited string key.""" + parts = [ + getattr(k, "key", str(k)) if not isinstance(k, (str, int)) else str(k) + for k in path_tuple + ] + if "layers_remainder" in parts: + num_scanned = 0 + if root_state is not None: + pure_s = root_state.to_pure_dict() if hasattr( + root_state, "to_pure_dict") else root_state + if isinstance(pure_s, dict): + + def _count_scanned(d): + if not isinstance(d, dict): + return 0 + if "scanned_blocks" in d and isinstance(d["scanned_blocks"], dict): + sb = d["scanned_blocks"] + if "layers" in sb and isinstance(sb["layers"], dict): + return len(sb["layers"]) + return len([ + k for k in sb.keys() + if k.startswith("layers_") or str(k).isdigit() + ]) + for v in d.values(): + if isinstance(v, dict): + c = _count_scanned(v) + if c > 0: + return c + return 0 + + num_scanned = _count_scanned(pure_s) + + rem_idx = parts.index("layers_remainder") + new_parts = list(parts[:rem_idx]) + for p in parts[rem_idx + 1:]: + if p == "layers": + continue + if p.startswith("layers_") and p[7:].isdigit(): + global_idx = num_scanned + int(p[7:]) + new_parts.append(f"layers_{global_idx}") + elif p.isdigit(): + global_idx = num_scanned + int(p) + new_parts.append(f"layers_{global_idx}") + else: + new_parts.append(p) + return _norm_path_key(new_parts) + + return _norm_path_key(parts) + + +def _update_nnx_state_from_pure_dict(nnx_state, pure_dict): + """Overlays pure dictionary parameters onto target NNX state variables.""" + if isinstance(nnx_state, dict): + target_state = nnx_state + elif hasattr(nnx_state, "to_pure_dict"): + target_state = nnx_state + else: + return + + pure_dict = pure_dict.to_pure_dict() if hasattr(pure_dict, + "to_pure_dict") else pure_dict + if not isinstance(pure_dict, (dict, nnx.State)): + return + + pure_dict = train_state_nnx._rename_nnx_to_linen_layers(pure_dict) # pylint: disable=protected-access + + # Convert pure_dict to a flat mapping of norm_path -> raw_array + flat_updates = {} + if isinstance(pure_dict, (dict, nnx.State)): + pure_flat = _flatten_and_norm_dict(pure_dict) + for k, v in pure_flat.items(): + raw_val = _rebuild_qwix_types(v) + if isinstance(raw_val, nnx.State): + continue + flat_updates[k] = raw_val + + leaves, _ = jax.tree_util.tree_flatten_with_path( + nnx_state, is_leaf=lambda *args: isinstance(args[-1], nnx.Variable)) + for path, var in leaves: + path_tuple = tuple(getattr(k, "key", str(k)) for k in path) + norm_p = _norm_nnx_path_key(path_tuple, target_state) + if norm_p in flat_updates: + val = flat_updates[norm_p] + if val is not None and not isinstance(val, jax.ShapeDtypeStruct): + _update_leaf_var(var, val, target_state, path) def _expected_and_restored_params(abstract_nnx_state, restored_linen): @@ -83,7 +364,12 @@ def _expected_and_restored_params(abstract_nnx_state, restored_linen): Splits the abstract by Variable type (nnx.Param) so only real weights are compared -- rngs/dropout/batch stats live in `nnx_aux` and are restored separately. """ - want = nnx.split_state(abstract_nnx_state, nnx.Param, ...)[0].to_pure_dict().get("model", {}) + lora_params = nnx.filter_state(abstract_nnx_state, nnx.LoRAParam) + if lora_params: + want = lora_params.to_pure_dict().get("model", {}) + else: + want = nnx.split_state(abstract_nnx_state, nnx.Param, + ...)[0].to_pure_dict().get("model", {}) have = restored_linen.get("params", {}).get("params", {}) return want, have @@ -117,16 +403,18 @@ def _linen_items_to_nnx(restored_linen, abstract_nnx_state): caller's abstract is untouched. Leaves the checkpoint didn't carry -- including the caches it never stores -- stay unmaterialized `ShapeDtypeStruct`s; the caller fills them from a fresh init. """ - linen_state, aux_state, ephemeral = train_state_nnx.split_for_checkpoint(abstract_nnx_state) + linen_state, aux_state, ephemeral = train_state_nnx.split_for_checkpoint( + abstract_nnx_state) weights = train_state_nnx.from_linen_checkpoint_dict(restored_linen) if "model" in weights: - nnx.replace_by_pure_dict(linen_state, {"model": weights["model"]}) + _update_nnx_state_from_pure_dict(linen_state, {"model": weights["model"]}) if "optimizer" in weights: - nnx.replace_by_pure_dict(linen_state, {"optimizer": weights["optimizer"]}) + _update_nnx_state_from_pure_dict(linen_state, + {"optimizer": weights["optimizer"]}) nnx_aux = restored_linen.get("nnx_aux") if nnx_aux: - nnx.replace_by_pure_dict(aux_state, nnx_aux) + _update_nnx_state_from_pure_dict(aux_state, nnx_aux) return nnx.merge_state(linen_state, aux_state, ephemeral) @@ -152,12 +440,14 @@ def _load_linen_checkpoint_into_nnx( save_concurrent_gb=checkpoint_storage_concurrent_gb, use_ocdbt=use_ocdbt, use_zarr3=use_zarr3, - ) - ) + )) restore_args = ocp.checkpoint_utils.construct_restore_args(linen_abstract) - restored = ocp.args.PyTreeRestore(item=linen_abstract, restore_args=restore_args, partial_restore=True) + restored = ocp.args.PyTreeRestore(item=linen_abstract, + restore_args=restore_args, + partial_restore=True) restored = ckptr.restore(epath.Path(path), args=restored) - _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_nnx_state, restored)) + _raise_on_weight_mismatch( + *_expected_and_restored_params(abstract_nnx_state, restored)) return _linen_items_to_nnx(restored, abstract_nnx_state) @@ -173,7 +463,9 @@ def _restore_emergency_linen_checkpoint_into_nnx( carries it too; it's restored when present and otherwise kept at its fresh init value. A genuinely-missing weight raises. """ - max_logging.log(f"Restoring emergency Linen-layout checkpoint into NNX state at step {step}") + max_logging.log( + f"Restoring emergency Linen-layout checkpoint into NNX state at step {step}" + ) linen_abstract = train_state_nnx.to_checkpoint_dict(abstract_nnx_state) restore_args = jax.tree_util.tree_map(map_to_pspec, linen_abstract) checkpoint_args = ocp.args.PyTreeRestore( @@ -181,8 +473,10 @@ def _restore_emergency_linen_checkpoint_into_nnx( restore_args=restore_args, partial_restore=True, ) - restored = checkpoint_manager.restore(step, args=Composite(state=checkpoint_args)).state - _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_nnx_state, restored)) + restored = checkpoint_manager.restore( + step, args=Composite(state=checkpoint_args)).state + _raise_on_weight_mismatch( + *_expected_and_restored_params(abstract_nnx_state, restored)) return _linen_items_to_nnx(restored, abstract_nnx_state) @@ -220,27 +514,34 @@ def _load_full_state_from_path( # pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state. if isinstance(abstract_unboxed_pre_state, nnx.State): return _load_linen_checkpoint_into_nnx( - path, abstract_unboxed_pre_state, checkpoint_storage_concurrent_gb, use_ocdbt, use_zarr3 - ) - context = ocp_v1.Context(checkpoint_layout=ocp_v1.options.CheckpointLayout.ORBAX) + path, abstract_unboxed_pre_state, checkpoint_storage_concurrent_gb, + use_ocdbt, use_zarr3) + context = ocp_v1.Context( + checkpoint_layout=ocp_v1.options.CheckpointLayout.ORBAX) with context: return ocp_v1.load_pytree(path, abstract_unboxed_pre_state) elif source_checkpoint_layout == "safetensors": - context = ocp_v1.Context(checkpoint_layout=ocp_v1.options.CheckpointLayout.SAFETENSORS) + context = ocp_v1.Context( + checkpoint_layout=ocp_v1.options.CheckpointLayout.SAFETENSORS) with context: metadata = ocp_v1.pytree_metadata(path) simple_abstract_state = metadata.metadata - shardings = sharding_utils.construct_maximal_shardings(simple_abstract_state) + shardings = sharding_utils.construct_maximal_shardings( + simple_abstract_state) def combine_sharding(sds, shardings): - return jax.ShapeDtypeStruct(shape=sds.shape, dtype=sds.dtype, sharding=shardings) + return jax.ShapeDtypeStruct(shape=sds.shape, + dtype=sds.dtype, + sharding=shardings) - sharded_abstract_state = jax.tree.map(combine_sharding, simple_abstract_state, shardings) + sharded_abstract_state = jax.tree.map(combine_sharding, + simple_abstract_state, shardings) pre_transformed_state = ocp_v1.load_pytree(path, sharded_abstract_state) state = checkpoint_conversion_fn(pre_transformed_state) return state else: - raise ocp_v1.errors.InvalidLayoutError(f"Unknown checkpoint layout: {source_checkpoint_layout}") + raise ocp_v1.errors.InvalidLayoutError( + f"Unknown checkpoint layout: {source_checkpoint_layout}") else: # pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state. if isinstance(abstract_unboxed_pre_state, nnx.State): @@ -267,7 +568,9 @@ def combine_sharding(sds, shardings): lambda x: ocp.type_handlers.ArrayRestoreArgs(sharding=x.sharding), restore_target, ) - return ocp.Checkpointer(handler).restore(p, restore_target, restore_args=restore_args) + return ocp.Checkpointer(handler).restore(p, + restore_target, + restore_args=restore_args) def create_orbax_checkpoint_manager( @@ -294,47 +597,51 @@ def create_orbax_checkpoint_manager( max_logging.log("Checkpointing disabled, not creating checkpoint manager.") return None - max_logging.log(f"Creating checkpoint manager with ocdbt={use_ocdbt} and zarr3={use_zarr3}") + max_logging.log( + f"Creating checkpoint manager with ocdbt={use_ocdbt} and zarr3={use_zarr3}" + ) # Base configuration for all dataset types item_names = ("items",) # we need to use ocdbt and zarr3 to control max file size in the checkpoint item_handlers = { - "items": PyTreeCheckpointHandler( - restore_concurrent_gb=checkpoint_storage_concurrent_gb, - save_concurrent_gb=checkpoint_storage_concurrent_gb, - use_ocdbt=use_ocdbt, - use_zarr3=use_zarr3, - ) + "items": + PyTreeCheckpointHandler( + restore_concurrent_gb=checkpoint_storage_concurrent_gb, + save_concurrent_gb=checkpoint_storage_concurrent_gb, + use_ocdbt=use_ocdbt, + use_zarr3=use_zarr3, + ) } if dataset_type is not None and dataset_type == "grain": item_names += ("iter",) - item_handlers["iter"] = grain_utility.GrainCheckpointHandler() # pyrefly: ignore[bad-assignment] + item_handlers["iter"] = grain_utility.GrainCheckpointHandler( + ) # pyrefly: ignore[bad-assignment] # local storage checkpoint needs parent directory created p = gcs_utils.mkdir_and_check_permissions(checkpoint_dir) if enable_continuous_checkpointing: max_logging.log("Enabling policy for continuous checkpointing.") - save_decision_policy = save_decision_policy_lib.ContinuousCheckpointingPolicy() + save_decision_policy = save_decision_policy_lib.ContinuousCheckpointingPolicy( + ) elif enable_autocheckpoint: max_logging.log("Enabling policy for autocheckpoint.") - save_decision_policy = save_decision_policy_lib.AnySavePolicy( - [ - save_decision_policy_lib.PreemptionCheckpointingPolicy(), - save_decision_policy_lib.FixedIntervalPolicy(save_interval_steps), - ] - ) + save_decision_policy = save_decision_policy_lib.AnySavePolicy([ + save_decision_policy_lib.PreemptionCheckpointingPolicy(), + save_decision_policy_lib.FixedIntervalPolicy(save_interval_steps), + ]) else: max_logging.log("Enabling policy for fixed interval checkpointing.") - save_decision_policy = save_decision_policy_lib.FixedIntervalPolicy(interval=save_interval_steps) - preservation_policy = preservation_policy_lib.LatestN(max_num_checkpoints_to_keep) + save_decision_policy = save_decision_policy_lib.FixedIntervalPolicy( + interval=save_interval_steps) + preservation_policy = preservation_policy_lib.LatestN( + max_num_checkpoints_to_keep) async_options = None if enable_continuous_checkpointing: - async_options = ocp.AsyncOptions( - timeout_secs=int(datetime.timedelta(minutes=60).total_seconds()), - ) + async_options = ocp.AsyncOptions(timeout_secs=int( + datetime.timedelta(minutes=60).total_seconds()),) manager = ocp.CheckpointManager( p, item_names=item_names, @@ -364,7 +671,8 @@ def print_save_message(step, async_checkpointing): def load_state_if_possible( checkpoint_manager: CheckpointManager | None, - data_iterator: MultiHostDataLoadIterator | list[MultiHostDataLoadIterator] | None, + data_iterator: MultiHostDataLoadIterator | list[MultiHostDataLoadIterator] | + None, load_parameters_from_path: str, load_full_state_from_path: str, checkpoint_storage_concurrent_gb: int, @@ -407,9 +715,12 @@ def load_state_if_possible( """ if checkpoint_manager is not None: - max_logging.log("checkpoint manager exists so trying to load this run's existing checkpoint") + max_logging.log( + "checkpoint manager exists so trying to load this run's existing checkpoint" + ) - step = checkpoint_manager.latest_step() if step < 0 else step # pyrefly: ignore[bad-assignment] + step = checkpoint_manager.latest_step( + ) if step < 0 else step # pyrefly: ignore[bad-assignment] if step is not None: max_logging.log(f"restoring from this run's directory step {step}") @@ -419,9 +730,11 @@ def map_to_pspec(data): pspec = data.sharding.spec mesh = data.sharding.mesh replica_axis_index = 0 - replica_devices = grain_utility.replica_devices(mesh.devices, replica_axis_index) + replica_devices = grain_utility.replica_devices(mesh.devices, + replica_axis_index) replica_mesh = jax.sharding.Mesh(replica_devices, mesh.axis_names) - single_replica_sharding = jax.sharding.NamedSharding(replica_mesh, pspec) + single_replica_sharding = jax.sharding.NamedSharding( + replica_mesh, pspec) return ocp.type_handlers.SingleReplicaArrayRestoreArgs( sharding=jax.sharding.NamedSharding(mesh, pspec), @@ -435,7 +748,9 @@ def map_to_pspec(data): replica_axis_index=0, broadcast_memory_limit_bytes=1024 * 1024 * 1000, # 1000 MB limit ) - ocp.type_handlers.register_type_handler(jax.Array, array_handler, override=True) + ocp.type_handlers.register_type_handler(jax.Array, + array_handler, + override=True) # pure_nnx saves in the Linen on-disk layout; restore that layout (weights + # opt_state + step + nnx_aux), restoring the grain iterator in place when @@ -445,22 +760,25 @@ def map_to_pspec(data): checkpoint_manager, (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager), ): - linen_abstract = train_state_nnx.to_checkpoint_dict(abstract_unboxed_pre_state) + linen_abstract = train_state_nnx.to_checkpoint_dict( + abstract_unboxed_pre_state) restore_args = jax.tree_util.tree_map(map_to_pspec, linen_abstract) - checkpoint_args = ocp.args.PyTreeRestore(item=linen_abstract, restore_args=restore_args, partial_restore=True) - if ( - dataset_type == "grain" - and data_iterator - and not isinstance(data_iterator, PlaceHolderDataIterator) - and (checkpoint_manager.directory / str(step) / "iter").exists() - ): + checkpoint_args = ocp.args.PyTreeRestore(item=linen_abstract, + restore_args=restore_args, + partial_restore=True) + if (dataset_type == "grain" and data_iterator and + not isinstance(data_iterator, PlaceHolderDataIterator) and + (checkpoint_manager.directory / str(step) / "iter").exists()): restored, _ = grain_utility.restore_grain_iterator( - checkpoint_manager, step, data_iterator, checkpoint_args, expansion_factor_real_data - ) + checkpoint_manager, step, data_iterator, checkpoint_args, + expansion_factor_real_data) else: - restored = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args)) - _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_unboxed_pre_state, restored["items"])) - restored_nnx = _linen_items_to_nnx(restored["items"], abstract_unboxed_pre_state) + restored = checkpoint_manager.restore( + step, args=Composite(items=checkpoint_args)) + _raise_on_weight_mismatch(*_expected_and_restored_params( + abstract_unboxed_pre_state, restored["items"])) + restored_nnx = _linen_items_to_nnx(restored["items"], + abstract_unboxed_pre_state) return ({"items": restored_nnx}, None) if isinstance(abstract_unboxed_pre_state, nnx.State) and isinstance( @@ -488,9 +806,9 @@ def map_to_pspec(data): ) match (checkpoint_manager, dataset_type, data_iterator): - # Case 1: Matches if 'checkpoint_manager' is an instance of either EmergencyCheckpointManager - # or EmergencyReplicatorCheckpointManager. The '_' indicates that 'dataset_type' and - # 'data_iterator' can be any value and aren't used in this pattern. + # Case 1: Matches if 'checkpoint_manager' is an instance of either EmergencyCheckpointManager + # or EmergencyReplicatorCheckpointManager. The '_' indicates that 'dataset_type' and + # 'data_iterator' can be any value and aren't used in this pattern. case (checkpoint_manager, _, _) if isinstance( checkpoint_manager, ( @@ -498,7 +816,8 @@ def map_to_pspec(data): EmergencyReplicatorCheckpointManager, ), ): - restored = checkpoint_manager.restore(step, args=Composite(state=checkpoint_args)).state + restored = checkpoint_manager.restore( + step, args=Composite(state=checkpoint_args)).state return ( restored, None, @@ -509,12 +828,9 @@ def map_to_pspec(data): checkpoint_manager, dataset_type, data_iterator, - ) if ( - dataset_type == "grain" - and data_iterator - and not isinstance(data_iterator, PlaceHolderDataIterator) - and (checkpoint_manager.directory / str(step) / "iter").exists() - ): + ) if (dataset_type == "grain" and data_iterator and + not isinstance(data_iterator, PlaceHolderDataIterator) and + (checkpoint_manager.directory / str(step) / "iter").exists()): return grain_utility.restore_grain_iterator( checkpoint_manager, step, @@ -525,17 +841,22 @@ def map_to_pspec(data): # Case 3: Default/Fallback case. # This case acts as a wildcard ('_') and matches if none of the preceding cases were met. case _: - restored = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args)) + restored = checkpoint_manager.restore( + step, args=Composite(items=checkpoint_args)) return (restored, None) if source_checkpoint_layout == "safetensors_dynamic": path = load_parameters_from_path or load_full_state_from_path - max_logging.log(f"Dynamic On-the-Fly Formatting: Loading SafeTensors from {path}") + max_logging.log( + f"Dynamic On-the-Fly Formatting: Loading SafeTensors from {path}") - return load_safetensors_dynamic_state(path, abstract_unboxed_pre_state, maxtext_config) + return load_safetensors_dynamic_state(path, abstract_unboxed_pre_state, + maxtext_config) elif load_parameters_from_path != "": if isinstance(abstract_unboxed_pre_state, nnx.State): - _, params, _ = nnx.split(abstract_unboxed_pre_state.model, nnx.Param, ...) + params = (abstract_unboxed_pre_state.model if hasattr( + abstract_unboxed_pre_state, "model") else + abstract_unboxed_pre_state["model"]) else: params = abstract_unboxed_pre_state.params @@ -548,7 +869,8 @@ def map_to_pspec(data): ) return None, restored_params elif load_full_state_from_path != "": - max_logging.log(f"Loading full state from path: {load_full_state_from_path}") + max_logging.log( + f"Loading full state from path: {load_full_state_from_path}") restored_state = _load_full_state_from_path( path=load_full_state_from_path, abstract_unboxed_pre_state=abstract_unboxed_pre_state, @@ -577,8 +899,8 @@ def setup_checkpoint_logger(config) -> Any | None: # pytype: disable=attribute- if config.enable_checkpoint_cloud_logger: logger_name = f"goodput_{config.run_name}" orbax_cloud_logger = ocp.logging.CloudLogger( - options=ocp.logging.CloudLoggerOptions(job_name=config.run_name, logger_name=logger_name) - ) + options=ocp.logging.CloudLoggerOptions(job_name=config.run_name, + logger_name=logger_name)) max_logging.log("Successfully set up checkpoint cloud logger.") return orbax_cloud_logger @@ -598,49 +920,83 @@ def load_params_from_path( # On disk the weights live at `params/params/...`: an outer key naming the item, and Flax's # `params` collection inside it. A Linen TrainState.params is that collection; an NNX params # state sits one level below it (bare weights), so wrap it going in and unwrap it coming out. - is_nnx = isinstance(abstract_unboxed_params, nnx.State) - want = abstract_unboxed_params.to_pure_dict() if is_nnx else abstract_unboxed_params - params_collection = {"params": want} if is_nnx else want - - # *_concurrent_gb should be set for large models, the default is 96. - max_logging.log(f"Creating checkpoint manager with ocdbt={use_ocdbt} and zarr3={use_zarr3}") - ckptr = ocp.Checkpointer( - ocp.PyTreeCheckpointHandler( - restore_concurrent_gb=checkpoint_storage_concurrent_gb, - save_concurrent_gb=checkpoint_storage_concurrent_gb, - use_ocdbt=use_ocdbt, - use_zarr3=use_zarr3, - ) - ) - - # This is a memory optimization. We don't want to restore the entire checkpoint - only the params. - # Rather than pass the entire abstract state, which could unnecessarily restore opt_state and such and waste - # memory, we instead specify here that we are just restoring the params field of the checkpoint - # (which itself may be a dictionary containing a key named 'params'). - restore_args = ocp.checkpoint_utils.construct_restore_args(params_collection) - restored = ckptr.restore( - epath.Path(load_parameters_from_path), - item={"params": params_collection}, - transforms={}, - restore_args={"params": restore_args}, - ) - restored_collection = restored["params"] - # `transforms={}` lets Orbax return an unmaterialized leaf for a weight the checkpoint lacks, - # and a stored array at its own shape rather than the target's. Either reaches the model and - # fails much later without naming the weight, so check here -- the params-only load - # (load_parameters_path, e.g. SFT) has no init state to fall back on. - _raise_on_weight_mismatch(want, restored_collection["params"] if is_nnx else restored_collection) + is_nnx = isinstance(abstract_unboxed_params, (nnx.State, nnx.Module)) if is_nnx: - nnx.replace_by_pure_dict(abstract_unboxed_params, restored_collection["params"]) - return abstract_unboxed_params + unquant = lora_utils.restore_qlora_base_weights(abstract_unboxed_params) + pure_want = unquant.to_pure_dict() if hasattr(unquant, + "to_pure_dict") else unquant + candidate_wants = [ + pure_want, + train_state_nnx._rename_nnx_to_linen_layers(pure_want) + ] # pylint: disable=protected-access + max_logging.log( + f"Creating checkpoint manager with ocdbt={use_ocdbt} and zarr3={use_zarr3}" + ) + ckptr = ocp.Checkpointer( + ocp.PyTreeCheckpointHandler( + restore_concurrent_gb=checkpoint_storage_concurrent_gb, + save_concurrent_gb=checkpoint_storage_concurrent_gb, + use_ocdbt=use_ocdbt, + use_zarr3=use_zarr3, + )) + last_err = None + for want in candidate_wants: + params_collection = {"params": want} + restore_args = ocp.checkpoint_utils.construct_restore_args( + params_collection) + restored = ckptr.restore( + epath.Path(load_parameters_from_path), + item={"params": params_collection}, + transforms={}, + restore_args={"params": restore_args}, + ) + restored_collection = restored["params"] + try: + _raise_on_weight_mismatch(want, restored_collection["params"]) + max_logging.log( + "load_params_from_path: successfully restored parameters.") + _update_nnx_state_from_pure_dict(abstract_unboxed_params, + restored_collection["params"]) + return abstract_unboxed_params + except ValueError as e: + last_err = e + if last_err is not None: + raise last_err + else: + want = abstract_unboxed_params + params_collection = want + max_logging.log( + f"Creating checkpoint manager with ocdbt={use_ocdbt} and zarr3={use_zarr3}" + ) + ckptr = ocp.Checkpointer( + ocp.PyTreeCheckpointHandler( + restore_concurrent_gb=checkpoint_storage_concurrent_gb, + save_concurrent_gb=checkpoint_storage_concurrent_gb, + use_ocdbt=use_ocdbt, + use_zarr3=use_zarr3, + )) + restore_args = ocp.checkpoint_utils.construct_restore_args( + params_collection) + restored = ckptr.restore( + epath.Path(load_parameters_from_path), + item={"params": params_collection}, + transforms={}, + restore_args={"params": restore_args}, + ) + restored_collection = restored["params"] + _raise_on_weight_mismatch(want, restored_collection) + return restored_collection return restored_collection def save_params_to_path(checkpoint_dir, params, use_ocdbt=True, use_zarr3=True): """Save decode params in checkpoint at specified path.""" assert checkpoint_dir, "checkpoint_dir is not defined." - print(f"Saving quantized params checkpoint with use_ocdbt = {use_ocdbt} and use_zarr3 = {use_zarr3}") - orbax_checkpointer = ocp.PyTreeCheckpointer(use_ocdbt=use_ocdbt, use_zarr3=use_zarr3) + print( + f"Saving quantized params checkpoint with use_ocdbt = {use_ocdbt} and use_zarr3 = {use_zarr3}" + ) + orbax_checkpointer = ocp.PyTreeCheckpointer(use_ocdbt=use_ocdbt, + use_zarr3=use_zarr3) orbax_checkpointer.save(checkpoint_dir, {"params": params}, force=True) print(f"Quantized params checkpoint saved at: {checkpoint_dir}") @@ -677,12 +1033,15 @@ def _should_save_checkpoint_at_step(checkpoint_manager, step, config, force): base_checkpoint_due = bool(checkpoint_manager.should_save(step)) else: base_checkpoint_due = step % config.checkpoint_period == 0 - local_checkpoint_due = _uses_local_checkpoint_period(config) and step % config.local_checkpoint_period == 0 - autocheckpoint_due = config.enable_autocheckpoint and checkpoint_manager.reached_preemption(step) + local_checkpoint_due = _uses_local_checkpoint_period( + config) and step % config.local_checkpoint_period == 0 + autocheckpoint_due = config.enable_autocheckpoint and checkpoint_manager.reached_preemption( + step) return base_checkpoint_due or local_checkpoint_due or autocheckpoint_due -def _handle_post_checkpoint_preemption(checkpoint_manager, step, force_ckpt_save): +def _handle_post_checkpoint_preemption(checkpoint_manager, step, + force_ckpt_save): """Waits on final/preemption saves and raises if preempted.""" reached_preemption = checkpoint_manager.reached_preemption(step) if force_ckpt_save or reached_preemption: @@ -691,7 +1050,11 @@ def _handle_post_checkpoint_preemption(checkpoint_manager, step, force_ckpt_save raise exceptions.StopTraining("Job is preempted.") -def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step=None): +def maybe_save_checkpoint(checkpoint_manager, + state, + config, + data_iterator, + step=None): """Save checkpoint if checkpointing is enabled.""" if checkpoint_manager is None: return @@ -704,7 +1067,8 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step else: if config.pure_nnx: # Under DiLoCo the step lives on the DiLoCoTrainState; otherwise on the optimizer. - actual_step = int(state.step if config.enable_diloco else state.optimizer.step) - 1 + actual_step = int( + state.step if config.enable_diloco else state.optimizer.step) - 1 else: # Linen TrainState has .step attribute actual_step = int(state.step) - 1 @@ -715,14 +1079,18 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step # without an explicit 'step' (implying it's a checkpoint save for final step), # AND the 'actual_step' is a valid step, # AND it's not a step that would normally trigger a checkpoint save. - force_ckpt_save = step is None and actual_step != -1 and (actual_step % config.checkpoint_period != 0) + force_ckpt_save = step is None and actual_step != -1 and ( + actual_step % config.checkpoint_period != 0) - if not _should_save_checkpoint_at_step(checkpoint_manager, actual_step, config, force_ckpt_save): - _handle_post_checkpoint_preemption(checkpoint_manager, actual_step, force_ckpt_save) + if not _should_save_checkpoint_at_step(checkpoint_manager, actual_step, + config, force_ckpt_save): + _handle_post_checkpoint_preemption(checkpoint_manager, actual_step, + force_ckpt_save) return if checkpoint_manager.latest_step() == actual_step: - max_logging.log(f"Checkpoint for step {actual_step} already exists, skipping save.") + max_logging.log( + f"Checkpoint for step {actual_step} already exists, skipping save.") return if config.pure_nnx: @@ -730,28 +1098,37 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step if config.enable_diloco: # DiLoCoTrainState: persist the synchronized global model (outer params). # The per-replica inner optimizer / outer-momentum state is not checkpointed. - step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step - state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}}) + step_value = state.step.get_value() if hasattr( + state.step, "get_value") else state.step + state = train_state_nnx.to_linen_checkpoint_dict({ + "model": state.params, + "optimizer": { + "step": step_value + } + }) else: # rngs/dropout/batch-stats are packed under items/nnx_aux so the RNG/dropout # stream continues across resumes instead of resetting to a base key. state = train_state_nnx.to_checkpoint_dict(state) try: - checkpoint_saved = save_checkpoint(checkpoint_manager, actual_step, state, config, data_iterator, force_ckpt_save) + checkpoint_saved = save_checkpoint(checkpoint_manager, actual_step, state, + config, data_iterator, force_ckpt_save) if checkpoint_saved: print_save_message(actual_step, config.async_checkpointing) if config.elastic_enabled: elastic_utils.maybe_elastic_scale_up(config, checkpoint_manager) except elastic_utils.manager.ScaleUpSignalError as e: if config.elastic_enabled: - max_logging.log(f"Elastic event detected, letting exception bubble up: {e}") + max_logging.log( + f"Elastic event detected, letting exception bubble up: {e}") raise else: raise exceptions.StopTraining("Job is preempted.") from e except jax.errors.JaxRuntimeError as e: if config.elastic_enabled: - max_logging.log(f"Elastic event detected, letting exception bubble up: {e}") + max_logging.log( + f"Elastic event detected, letting exception bubble up: {e}") raise else: raise exceptions.StopTraining("Job is preempted.") from e @@ -760,18 +1137,24 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step # Wait for any pending checkpoint save to finish during preemption or final # step save, then raise upon preemption. - _handle_post_checkpoint_preemption(checkpoint_manager, actual_step, force_ckpt_save) + _handle_post_checkpoint_preemption(checkpoint_manager, actual_step, + force_ckpt_save) -def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=None, force=False): +def save_checkpoint(checkpoint_manager, + step, + state, + config=None, + data_iterator=None, + force=False): """Wrapper for saving checkpoint.""" if config and config.enable_checkpointing: - if ( - force - or (step % config.checkpoint_period == 0 and not config.enable_continuous_checkpointing) - or (_uses_local_checkpoint_period(config) and step % config.local_checkpoint_period == 0) - or (config.enable_autocheckpoint and checkpoint_manager.reached_preemption(step)) - ): + if (force or (step % config.checkpoint_period == 0 and + not config.enable_continuous_checkpointing) or + (_uses_local_checkpoint_period(config) and + step % config.local_checkpoint_period == 0) or + (config.enable_autocheckpoint and + checkpoint_manager.reached_preemption(step))): blocking_until_ready_start = time.time() max_logging.log(f"Waiting for step {step} to finish before checkpoint...") # We block here on the step finishing so that our checkpointing metrics @@ -779,34 +1162,33 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator= jax.block_until_ready(state) max_logging.log( f"Waited {time.time() - blocking_until_ready_start} seconds for step " - f"{step} to finish before starting checkpointing." - ) + f"{step} to finish before starting checkpointing.") # specify chunk_byte_size to force orbax to control maximum file size in checkpoint - chunk_byte_size = ( - config.checkpoint_storage_target_data_file_size_bytes if config else DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE - ) + chunk_byte_size = (config.checkpoint_storage_target_data_file_size_bytes + if config else DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE) checkpoint_args = ocp.args.PyTreeSave( item=state, - save_args=jax.tree.map(lambda _: ocp.SaveArgs(chunk_byte_size=chunk_byte_size), state), + save_args=jax.tree.map( + lambda _: ocp.SaveArgs(chunk_byte_size=chunk_byte_size), state), ocdbt_target_data_file_size=chunk_byte_size, ) + save_args_composite = {"items": checkpoint_args} - if config and config.dataset_type == "grain" and not isinstance(data_iterator, PlaceHolderDataIterator): + if config and config.dataset_type == "grain" and not isinstance( + data_iterator, PlaceHolderDataIterator): if isinstance(data_iterator, RemoteIteratorWrapper): # Pass the wrapper directly; GrainCheckpointHandler will call save_state with the step save_args_composite["iter"] = grain_utility.GrainCheckpointSave( - item=data_iterator - ) # pyrefly: ignore[bad-assignment] + item=data_iterator) # pyrefly: ignore[bad-assignment] elif not isinstance(data_iterator, list) and isinstance( - data_iterator.local_iterator, ElasticIterator - ): # pyrefly: ignore[missing-attribute] + data_iterator.local_iterator, + ElasticIterator): # pyrefly: ignore[missing-attribute] # ElasticIterator checkpoints a single global scalar shared by all shards. save_args_composite["iter"] = grain_utility.GrainCheckpointSave( - item=data_iterator.local_iterator - ) # pyrefly: ignore[bad-assignment] + item=data_iterator.local_iterator) # pyrefly: ignore[bad-assignment] else: if not isinstance(data_iterator, list): data_iterator = [data_iterator] @@ -817,26 +1199,29 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator= for i, data_iter in enumerate(data_iterator): process_index = jax.process_index() + i * jax.process_count() grain_iters_to_save.append( - (data_iter.local_iterator, process_index, process_count_total) - ) # pyrefly: ignore[missing-attribute] + (data_iter.local_iterator, process_index, + process_count_total)) # pyrefly: ignore[missing-attribute] save_args_composite["iter"] = grain_utility.GrainCheckpointSave( - item=grain_iters_to_save - ) # pyrefly: ignore[bad-assignment] + item=grain_iters_to_save) # pyrefly: ignore[bad-assignment] custom_metadata = {} if config: if hasattr(config, "scan_layers"): custom_metadata["scan_layers"] = config.scan_layers - if hasattr(config, "lora") and config.lora and getattr(config.lora, "lora_rank", 0) > 0: + if hasattr(config, "lora") and config.lora and getattr( + config.lora, "lora_rank", 0) > 0: custom_metadata["lora"] = config.lora.model_dump() match (checkpoint_manager, config, data_iterator): case (checkpoint_manager, _, _) if isinstance( - checkpoint_manager, (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager) - ): + checkpoint_manager, + (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager)): emergency_checkpointing.replicator_error_handler(config) - return checkpoint_manager.save(step, args=Composite(state=checkpoint_args), force=force) + return checkpoint_manager.save(step, + args=Composite(state=checkpoint_args), + force=force) case _: - return checkpoint_manager.save( - step, args=Composite(**save_args_composite), force=force, custom_metadata=custom_metadata - ) + return checkpoint_manager.save(step, + args=Composite(**save_args_composite), + force=force, + custom_metadata=custom_metadata) diff --git a/src/maxtext/common/train_state_nnx.py b/src/maxtext/common/train_state_nnx.py index 7d73f45d4a..d4efa6ee34 100644 --- a/src/maxtext/common/train_state_nnx.py +++ b/src/maxtext/common/train_state_nnx.py @@ -11,7 +11,6 @@ # 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. - """The NNX Unified TrainState.""" from typing import Any @@ -51,10 +50,13 @@ def apply_gradients(self, grads: Any, **kwargs): raise RuntimeError( "Cannot call apply_gradients on a TrainStateNNX initialized without" " an optimizer. This usually happens when the state was created for" - " inference only." - ) + " inference only.") self.optimizer.update(self.model, grads, **kwargs) + def to_pure_dict(self): + """Returns the pure dict representation of the NNX state.""" + return nnx.state(self).to_pure_dict() + # On-disk checkpoint format. # @@ -81,7 +83,9 @@ def _cast_step(step, dtype): values. """ if isinstance(step, jax.ShapeDtypeStruct): - return jax.ShapeDtypeStruct(step.shape, dtype, sharding=getattr(step, "sharding", None)) + return jax.ShapeDtypeStruct(step.shape, + dtype, + sharding=getattr(step, "sharding", None)) return jnp.asarray(step, dtype=dtype) @@ -108,7 +112,12 @@ def _wrap_mu_nu_with_params(state): """Wraps mu/nu under an inner 'params' key (the Linen collection).""" if not isinstance(state, dict): return state - return {k: {"params": v} if k in ("mu", "nu") and isinstance(v, dict) else v for k, v in state.items()} + return { + k: { + "params": v + } if k in ("mu", "nu") and isinstance(v, dict) else + v for k, v in state.items() + } def _as_chain_index(key): @@ -137,13 +146,80 @@ def _opt_state_to_linen(opt_state): return chain +def _rename_nnx_to_linen_layers(d): + """Converts NNX nested layer dicts ('scanned_blocks/layers/0' or 'layers_remainder/layers/0') + to Linen-style flat layer keys ('layers_0', 'layers_1') for on-disk checkpoint compatibility. + """ + if not isinstance(d, dict): + return d + res = {} + num_scanned = 0 + + def _count_scanned(d_sub): + if not isinstance(d_sub, dict): + return 0 + if "scanned_blocks" in d_sub and isinstance(d_sub["scanned_blocks"], dict): + sb = d_sub["scanned_blocks"] + if "layers" in sb and isinstance(sb["layers"], dict): + return len(sb["layers"]) + return len( + [k for k in sb.keys() if k.startswith("layers_") or str(k).isdigit()]) + for v_sub in d_sub.values(): + if isinstance(v_sub, dict): + c = _count_scanned(v_sub) + if c > 0: + return c + return 0 + + num_scanned = _count_scanned(d) + + for k, v in d.items(): + if k == "scanned_blocks" and isinstance(v, dict): + if "layers" in v and isinstance(v["layers"], dict): + for l_idx, l_val in v["layers"].items(): + layer_key = f"layers_{l_idx}" if str(l_idx).isdigit() else str(l_idx) + res[layer_key] = _rename_nnx_to_linen_layers(l_val) + else: + for sub_k, sub_v in v.items(): + layer_key = sub_k if sub_k.startswith("layers_") else ( + f"layers_{sub_k}" if str(sub_k).isdigit() else sub_k) + res[layer_key] = _rename_nnx_to_linen_layers(sub_v) + elif k == "layers_remainder" and isinstance(v, dict): + if "layers" in v and isinstance(v["layers"], dict): + for l_idx, l_val in v["layers"].items(): + idx_int = int(l_idx) if str(l_idx).isdigit() else 0 + global_idx = num_scanned + idx_int + layer_key = f"layers_{global_idx}" + res[layer_key] = _rename_nnx_to_linen_layers(l_val) + else: + for sub_k, sub_v in v.items(): + if sub_k.startswith("layers_") and sub_k[7:].isdigit(): + idx_int = int(sub_k[7:]) + elif str(sub_k).isdigit(): + idx_int = int(sub_k) + else: + idx_int = 0 + global_idx = num_scanned + idx_int + layer_key = f"layers_{global_idx}" + res[layer_key] = _rename_nnx_to_linen_layers(sub_v) + else: + res[k] = _rename_nnx_to_linen_layers(v) + return res + + def to_linen_checkpoint_dict(nnx_pure_dict): """Reshapes a TrainStateNNX pure dict ({model, optimizer}) into the Linen on-disk layout.""" if not isinstance(nnx_pure_dict, dict): return nnx_pure_dict result = {} if "model" in nnx_pure_dict: - result["params"] = {"params": _strip_rng_state(nnx_pure_dict["model"])} + model_dict = _rename_nnx_to_linen_layers( + _strip_rng_state(nnx_pure_dict["model"])) + if isinstance(model_dict, + dict) and "params" in model_dict and len(model_dict) == 1: + result["params"] = model_dict + else: + result["params"] = {"params": model_dict} optimizer = nnx_pure_dict.get("optimizer") if isinstance(optimizer, dict): if "step" in optimizer: @@ -159,7 +235,8 @@ def _strip_mu_nu_params(state): if not isinstance(state, dict): return state return { - k: (v["params"] if k in ("mu", "nu") and isinstance(v, dict) and "params" in v else v) for k, v in state.items() + k: (v["params"] if k in ("mu", "nu") and isinstance(v, dict) and + "params" in v else v) for k, v in state.items() } @@ -171,7 +248,11 @@ def _opt_state_from_linen(opt_state): up with the model's opt_state. """ if isinstance(opt_state, list): - return {i: _strip_mu_nu_params(e) for i, e in enumerate(opt_state) if isinstance(e, dict)} + return { + i: _strip_mu_nu_params(e) + for i, e in enumerate(opt_state) + if isinstance(e, dict) + } if not isinstance(opt_state, dict): return opt_state return _strip_mu_nu_params(opt_state) @@ -202,43 +283,80 @@ def from_linen_checkpoint_dict(linen_pure_dict): def split_for_checkpoint(state: nnx.State): - """Partitions an nnx.State by its on-disk destination. - - Named by on-disk destination. Returns `(linen_state, aux, ephemeral)`: - linen_state: trainable weights (nnx.Param) and the optimizer -> Linen params/opt_state/step. - Its model side is pure nnx.Param, so it maps to the Linen `params` collection cleanly. - aux: rngs/dropout, batch stats, and any other persistent model variable (e.g. a - frozen routing table) -> nnx_aux. - ephemeral: caches and intermediates, which are recomputed and so are not checkpointed. - - The type-based catch-all `rest` mixes the optimizer (under "optimizer") with any custom - persistent variable (under "model"); they route to different places, so they're split apart - here. Grouping by the ephemeral types rather than whitelisting nnx.Param means a new Variable - subclass persists (via `aux`) instead of being silently dropped. Save and restore must - partition identically, so both go through here. + """Partition an NNX State into on-disk checkpoint collections. + + When LoRA parameters (nnx.LoRAParam) are present in the state, only LoRAParam + is saved in `params` (adapter-only saving), and frozen base parameters are + excluded to avoid duplicating base weights or saving quantized arrays. """ + has_lora = bool(nnx.filter_state(state, nnx.LoRAParam)) + flat = state.flat_state() + if not has_lora: + has_lora = any( + any( + isinstance(k, str) and ("lora_a" in k or "lora_b" in k) + for k in path) + for path, _ in flat) + + if has_lora and not bool(nnx.filter_state(state, nnx.LoRAParam)): + # Abstract/unboxed state where variable wrappers were stripped to ShapeDtypeStruct + params_flat = [] + custom_flat = [] + opt_flat = [] + aux_flat = [] + for path, val in flat: + if path and path[0] == "optimizer": + opt_flat.append((path, val)) + elif path and path[0] == "model": + if any(isinstance(k, str) and k in ("rngs", "dropout") for k in path): + aux_flat.append((path, val)) + elif any( + isinstance(k, str) and ("lora_a" in k or "lora_b" in k) + for k in path): + params_flat.append((path, val)) + else: + custom_flat.append((path, val)) + else: + aux_flat.append((path, val)) + params = nnx.State.from_flat_path(params_flat) + optimizer = nnx.State.from_flat_path(opt_flat) + custom = nnx.State.from_flat_path(custom_flat) + aux = nnx.State.from_flat_path(aux_flat) + linen_state = nnx.merge_state(params, optimizer) + aux_state = nnx.merge_state(aux, custom) + return linen_state, aux_state, nnx.State({}) + + param_type = nnx.LoRAParam if has_lora else nnx.Param + params, rng_state, batch_stats, caches, intermediates, rest = nnx.split_state( - state, nnx.Param, nnx.RngState, nnx.BatchStat, nnx.Cache, nnx.Intermediate, ... - ) - optimizer = nnx.State({"optimizer": rest["optimizer"]}) if "optimizer" in rest else nnx.State({}) - custom = nnx.State({"model": rest["model"]}) if "model" in rest else nnx.State({}) + state, param_type, nnx.RngState, nnx.BatchStat, nnx.Cache, + nnx.Intermediate, ...) + optimizer = nnx.State({"optimizer": rest["optimizer"] + }) if "optimizer" in rest else nnx.State({}) + custom = nnx.State({"model": rest["model"] + }) if ("model" in rest and not has_lora) else nnx.State({}) linen_state = nnx.merge_state(params, optimizer) + aux = nnx.merge_state(rng_state, batch_stats, custom) ephemeral = nnx.merge_state(caches, intermediates) return linen_state, aux, ephemeral -def to_checkpoint_dict(state: nnx.State): - """Reshapes an nnx.State into the on-disk checkpoint layout. +def to_checkpoint_dict(state: Any): + """Reshapes an nnx.State or TrainStateNNX into the on-disk checkpoint layout. - Weights (nnx.Param) map to the Linen `params` collection and the optimizer to + Weights (nnx.Param / nnx.LoRAParam) map to the Linen `params` collection and the optimizer to opt_state/step, so pure_nnx and Linen checkpoints stay interchangeable. Everything else that must persist -- rngs/dropout, batch stats, and any custom variable -- goes under an `nnx_aux` subtree. Works on a concrete state (save) or an abstract state (restore target). """ - linen_state, aux_state, _ = split_for_checkpoint(state) + nnx_state = state if isinstance(state, nnx.State) else nnx.state(state) + linen_state, aux_state, _ = split_for_checkpoint(nnx_state) pure = linen_state.to_pure_dict() - linen_dict = to_linen_checkpoint_dict({"model": pure.get("model", {}), "optimizer": pure.get("optimizer", {})}) + linen_dict = to_linen_checkpoint_dict({ + "model": pure.get("model", {}), + "optimizer": pure.get("optimizer", {}) + }) aux = aux_state.to_pure_dict() if aux: linen_dict["nnx_aux"] = aux diff --git a/src/maxtext/configs/post_train/lora_module_path.yml b/src/maxtext/configs/post_train/lora_module_path.yml index c84a7a5c8b..fab16e3b55 100644 --- a/src/maxtext/configs/post_train/lora_module_path.yml +++ b/src/maxtext/configs/post_train/lora_module_path.yml @@ -24,6 +24,7 @@ gemma3: "decoder/(scanned_blocks|layers_remainder|layers)/.*(self_attention/(que gemma4: "decoder/((scanned_blocks|layers_remainder)/)?layers.*/.*(self_attention/(query|key|value|out)|mlp/.*(wi_0|wi_1|wo|shared_experts/(wi_0|wi_1|wo)))" olmo3: "decoder/layers/.*(attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))" gpt3: "decoder/layers/(self_attention/(qkv_proj|out)|mlp/(wi|wo))" +gpt-oss: "decoder/layers/.*((GptOssAttention|self_attention)/(query|key|value|out)|(GptOssMlp|mlp)/(wi_0|wi_1|wo))" # Fallback for unverified models default: "decoder/layers/.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))" diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index ee92b93435..b9cce334db 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -25,7 +25,6 @@ from absl import app - import optax import pathwaysutils # pylint: disable=unused-import @@ -78,7 +77,8 @@ def get_first_step(model, state): if isinstance(model, nn.Module): return int(state.step) - if hasattr(state, "inner_state"): # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var + if hasattr(state, "inner_state" + ): # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var return int(state.step.get_value()) return int(state.optimizer.step.get_value()) @@ -88,7 +88,13 @@ def get_first_step(model, state): # ----------------------------------------------------------------------------- -def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_train=True): +def loss_fn(model, + config, + data, + dropout_rng, + params, + sparsity_state=None, + is_train=True): """loss_fn for both train and eval. Args: @@ -106,10 +112,10 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr # decimate proportion of data when per_device_batch_size<1 if is_train: for k, v in data.items(): - data[k] = v[: config.micro_batch_size_to_train_on, :] + data[k] = v[:config.micro_batch_size_to_train_on, :] else: for k, v in data.items(): - data[k] = v[: config.micro_batch_size_to_eval_on, :] + data[k] = v[:config.micro_batch_size_to_eval_on, :] mutable_collections = ["intermediates"] if config.mtp_num_layers > 0 and is_train: # The single model.apply call now triggers the entire chain if MTP is enabled: @@ -143,9 +149,13 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr data["inputs_position"], decoder_segment_ids=data["inputs_segmentation"], encoder_images=data["images"] if config.use_multimodal else None, - encoder_image_masks=data["image_masks"] if config.use_multimodal and "image_masks" in data else None, + encoder_image_masks=data["image_masks"] + if config.use_multimodal and "image_masks" in data else None, enable_dropout=config.enable_dropout if is_train else False, - rngs={"dropout": rng1, "params": aqt_rng}, # pyrefly: ignore[bad-argument-type] + rngs={ + "dropout": rng1, + "params": aqt_rng + }, # pyrefly: ignore[bad-argument-type] mutable=mutable_collections, decoder_target_tokens=data["targets"], decoder_target_mask=data["targets_segmentation"], @@ -158,11 +168,15 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr total_z_loss = 0.0 elif config.num_vocab_tiling > 1: hidden_state_key = ("intermediates", "decoder", "hidden_states") - hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0] - xent_sum, total_z_loss = vocab_tiling_linen_loss(hidden_states, data, config, model, params, is_train) + hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, + hidden_state_key)[0] + xent_sum, total_z_loss = vocab_tiling_linen_loss(hidden_states, data, + config, model, params, + is_train) else: one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size) - xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier) + xent, z_loss = max_utils.cross_entropy_with_logits( + logits, one_hot_targets, z_loss=config.z_loss_multiplier) xent = sharding.maybe_shard_with_logical( xent, @@ -192,7 +206,8 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr decoder_positions=data["inputs_position"], decoder_segment_ids=data["inputs_segmentation"], encoder_images=data["images"] if config.use_multimodal else None, - encoder_image_masks=data["image_masks"] if config.use_multimodal and "image_masks" in data else None, + encoder_image_masks=data["image_masks"] + if config.use_multimodal and "image_masks" in data else None, enable_dropout=config.enable_dropout if is_train else False, decoder_target_tokens=data["targets"], decoder_target_mask=data["targets_segmentation"], @@ -212,7 +227,8 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr # calculate_mtp_acceptance_rate find them at the same path as the Linen collections. if mtp_losses_state is not None and mtp_acceptance_state is not None: intermediate_outputs["mtp_losses"] = mtp_losses_state.to_pure_dict() - intermediate_outputs["mtp_acceptance"] = mtp_acceptance_state.to_pure_dict() + intermediate_outputs[ + "mtp_acceptance"] = mtp_acceptance_state.to_pure_dict() if (config.use_indexer and not config.indexer_sparse_training) and is_train: # In Dense Warm-up stage, we skip main model loss calculation for efficiency. @@ -221,11 +237,14 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr total_z_loss = 0.0 elif config.num_vocab_tiling > 1: hidden_state_key = ("decoder", "hidden_states") - hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0] - xent_sum, total_z_loss = vocab_tiling_nnx_loss(model, hidden_states, data, config, is_train) + hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, + hidden_state_key)[0] + xent_sum, total_z_loss = vocab_tiling_nnx_loss(model, hidden_states, data, + config, is_train) else: one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size) - xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier) + xent, z_loss = max_utils.cross_entropy_with_logits( + logits, one_hot_targets, z_loss=config.z_loss_multiplier) xent = sharding.maybe_shard_with_logical( xent, @@ -279,7 +298,8 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr # get indexer loss indexer_loss = 0.0 if config.use_indexer and config.indexer_loss_scaling_factor > 0.0: - indexer_losses = maxtext_utils.collect_intermediates_by_suffix(intermediate_outputs, "self_attention", "indexer_loss") + indexer_losses = maxtext_utils.collect_intermediates_by_suffix( + intermediate_outputs, "self_attention", "indexer_loss") if indexer_losses: indexer_loss = jnp.mean(jnp.concatenate(indexer_losses)) loss += indexer_loss @@ -289,7 +309,8 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr # get MoE load balance loss moe_lb_loss = 0.0 if config.num_experts > 1: - moe_lb_losses = maxtext_utils.collect_intermediates_by_suffix(intermediate_outputs, "moe_lb_loss") + moe_lb_losses = maxtext_utils.collect_intermediates_by_suffix( + intermediate_outputs, "moe_lb_loss") if moe_lb_losses: moe_lb_loss = jnp.mean(jnp.concatenate(moe_lb_losses)) loss += moe_lb_loss @@ -300,18 +321,19 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr moe_bias_updates = None if config.routed_bias and config.routed_bias_update_rate > 0.0: if isinstance(model, nn.Module): - nested_key = ("intermediates", "decoder", "moe_layers", "moe_bias_updates") - moe_bias_updates = maxtext_utils.get_nested_value(intermediate_outputs, nested_key, None) + nested_key = ("intermediates", "decoder", "moe_layers", + "moe_bias_updates") + moe_bias_updates = maxtext_utils.get_nested_value(intermediate_outputs, + nested_key, None) else: # NNX intermediates are model-rooted (no "intermediates" prefix), so match by # suffix instead. Unlike collect_intermediates_by_suffix we must not ravel: # the update is a 2-D matrix that's transposed at the apply site below. moe_bias_updates = next( - ( - val - for path, val in jax.tree_util.tree_leaves_with_path(intermediate_outputs) - if tuple(k.key for k in path if hasattr(k, "key"))[-1:] == ("moe_bias_updates",) - ), + (val for path, val in jax.tree_util.tree_leaves_with_path( + intermediate_outputs) + if tuple(k.key for k in path if hasattr(k, "key"))[-1:] == ( + "moe_bias_updates",)), None, ) if moe_bias_updates is not None: @@ -324,20 +346,35 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr intermediate_outputs["logits"] = logits aux = { - "intermediate_outputs": intermediate_outputs, - "xent_sum": xent_sum, - "z_loss": total_z_loss, - "total_weights": total_weights, - "moe_lb_loss": moe_lb_loss, - "indexer_loss": indexer_loss, - "moe_bias_updates": moe_bias_updates, - "mtp_loss": mtp_loss, - "batch_stats": (intermediate_outputs.get("batch_stats", None) if hasattr(intermediate_outputs, "get") else None), + "intermediate_outputs": + intermediate_outputs, + "xent_sum": + xent_sum, + "z_loss": + total_z_loss, + "total_weights": + total_weights, + "moe_lb_loss": + moe_lb_loss, + "indexer_loss": + indexer_loss, + "moe_bias_updates": + moe_bias_updates, + "mtp_loss": + mtp_loss, + "batch_stats": (intermediate_outputs.get("batch_stats", None) if hasattr( + intermediate_outputs, "get") else None), } return loss, aux -def train_step(model, config, state_mesh_shardings, params_shardings, state, data, dropout_rng=None): +def train_step(model, + config, + state_mesh_shardings, + params_shardings, + state, + data, + dropout_rng=None): """Training step for both Linen and NNX models. Args: @@ -376,12 +413,14 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat if isinstance(model, nn.Module): if config.shard_optimizer_over_data: params = jax.tree.map( - functools.partial(sharding.maybe_shard_with_name, shard_mode=config.shard_mode), + functools.partial(sharding.maybe_shard_with_name, + shard_mode=config.shard_mode), params, # pyrefly: ignore[unbound-name] params_shardings, ) sparsity_enabled = config.weight_sparsity_n and config.weight_sparsity_m - pure_params = params["params"] if sparsity_enabled else params # pyrefly: ignore[unbound-name] + pure_params = params[ + "params"] if sparsity_enabled else params # pyrefly: ignore[unbound-name] batch_stats = params.get("batch_stats", {}) grad_func = jax.value_and_grad(loss_fn, argnums=4, has_aux=True) @@ -395,9 +434,14 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat is_train=True, ) else: - owg_type = variablelib.variable_type_from_name("_overwrite_with_gradient", allow_register=True) + owg_type = variablelib.variable_type_from_name("_overwrite_with_gradient", + allow_register=True) custom_param_filter = nnx.Any(owg_type) - model_graphdef, curr_params, custom_params, rest = nnx.split(state.model, nnx.Param, custom_param_filter, ...) + train_param_type = (getattr(nnx, "LoRAParam", nnx.Param) if getattr( + getattr(config, "lora", None), "enable_lora", False) else nnx.Param) + nnx.pop(state.model, nnx.Intermediate) + model_graphdef, curr_params, custom_params, rest = nnx.split( + state.model, train_param_type, custom_param_filter, ...) if config.parameter_memory_host_offload: # Params are kept on host (pinned_host) in in_shardings. Move only Param # variables to device before the forward/backward pass so that all dot_general @@ -410,24 +454,51 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat is_leaf=lambda x: isinstance(x, NamedSharding), ) curr_params = jax.device_put(curr_params, device_param_shardings) - nnx.update(state.model, curr_params) # ensure state.model has device params for optimizer update + nnx.update(state.model, curr_params + ) # ensure state.model has device params for optimizer update if config.shard_optimizer_over_data: - curr_params = jax.tree.map( - functools.partial(sharding.maybe_shard_with_name, shard_mode=config.shard_mode), - curr_params, + param_sharding_lookup = {} + for p, s in jax.tree_util.tree_leaves_with_path( params_shardings, + is_leaf=lambda x: isinstance(x, (nnx.Variable, NamedSharding, jax. + sharding.Sharding))): + param_sharding_lookup[p] = s.get_value() if isinstance( + s, nnx.Variable) else s + + def _maybe_shard_param(path, var): + if path in param_sharding_lookup: + return sharding.maybe_shard_with_name(var, + param_sharding_lookup[path], + shard_mode=config.shard_mode) + return var + + curr_params = jax.tree_util.tree_map_with_path( + _maybe_shard_param, + curr_params, + is_leaf=lambda x: isinstance(x, nnx.Variable), ) nnx.update(state.model, curr_params) def diff_wrapper(curr_params, custom_params, rest, config, data): - local_model = nnx.merge(model_graphdef, curr_params, custom_params, rest, copy=True) - loss, aux = loss_fn(local_model, config, data, None, None, is_train=True) - _, _, _, new_rest = nnx.split(local_model, nnx.Param, custom_param_filter, ...) - return loss, (aux, new_rest) + local_model = nnx.merge(model_graphdef, + curr_params, + custom_params, + rest, + copy=True) + loss, aux = loss_fn(local_model, + config, + data, + None, + None, + is_train=True) + non_param_rest = nnx.state( + local_model, nnx.Not(nnx.Any(nnx.Param, nnx.Intermediate))) + return loss, (aux, non_param_rest) grad_func = jax.value_and_grad(diff_wrapper, argnums=(0, 1), has_aux=True) - (loss, (aux, new_rest)), (raw_grads, custom_grads) = grad_func(curr_params, custom_params, rest, config, data) - nnx.update(state.model, nnx.State.merge(custom_grads, new_rest)) + (loss, (aux, non_param_rest)), (raw_grads, custom_grads) = grad_func( + curr_params, custom_params, rest, config, data) + nnx.update(state.model, nnx.State.merge(custom_grads, non_param_rest)) raw_grads = jax.tree_util.tree_map( lambda x: x.astype(config.grad_dtype) if x.dtype == jnp.float32 else x, @@ -452,40 +523,39 @@ def diff_wrapper(curr_params, custom_params, rest, config, data): if isinstance(model, nn.Module): if config.gradient_clipping_threshold > 0: - grads = maxtext_utils.apply_gradient_clipping(raw_grads, state, config.gradient_clipping_threshold) + grads = maxtext_utils.apply_gradient_clipping( + raw_grads, state, config.gradient_clipping_threshold) else: grads = raw_grads if config.optimizer_memory_host_offload: - state = state.replace( - opt_state=jax.device_put( - state.opt_state, - jax.tree_util.tree_map( - lambda x: x.with_memory_kind(kind="device"), - state_mesh_shardings.opt_state, - ), - ) - ) + state = state.replace(opt_state=jax.device_put( + state.opt_state, + jax.tree_util.tree_map( + lambda x: x.with_memory_kind(kind="device"), + state_mesh_shardings.opt_state, + ), + )) # Move all parameters to device before optimizer update if config.parameter_memory_host_offload: - max_logging.log("\nMoving all parameters to device before optimizer update") + max_logging.log( + "\nMoving all parameters to device before optimizer update") def move(path, value): max_logging.log(f"train.py: Moving f{path} to device") return value.with_memory_kind(kind="device") - state = state.replace( - params=jax.device_put( - state.params, - jax.tree_util.tree_map_with_path(move, state_mesh_shardings.params), - ) - ) + state = state.replace(params=jax.device_put( + state.params, + jax.tree_util.tree_map_with_path(move, state_mesh_shardings.params), + )) # Re-wrap grads to match state.params structure if it's a dict of collections # (when weight_sparsity is enabled, params has both 'params' and 'batch_stats' keys). sparsity_enabled = config.weight_sparsity_n and config.weight_sparsity_m if sparsity_enabled: full_grads = {"params": grads} if "batch_stats" in state.params: - batch_stats_grads = jax.tree_util.tree_map(jnp.zeros_like, state.params.get("batch_stats", {})) + batch_stats_grads = jax.tree_util.tree_map( + jnp.zeros_like, state.params.get("batch_stats", {})) full_grads["batch_stats"] = batch_stats_grads full_grads = max_utils.unbox_logicallypartioned(full_grads) else: @@ -494,7 +564,11 @@ def move(path, value): if getattr(config, "skip_step_on_spikes", False): grad_norm = max_utils.l2norm_pytree(grads) # TrainState.apply_gradients doesn't pass **kwargs to tx.update, so we unpack it manually. - updates, new_opt_state = state.tx.update(grads, state.opt_state, state.params, loss=loss, grad_norm=grad_norm) + updates, new_opt_state = state.tx.update(grads, + state.opt_state, + state.params, + loss=loss, + grad_norm=grad_norm) new_params = optax.apply_updates(state.params, updates) new_state = state.replace( @@ -507,13 +581,16 @@ def move(path, value): # Apply updates for Auxiliary-Loss-Free load balancing for DeepSeek family if config.routed_bias and config.routed_bias_update_rate > 0.0 and moe_bias_updates is not None: - target_path = ("params", "decoder", "moe_layers", "DeepSeekMoeBlock_0", "MoeBlock_0", "gate", "bias") + target_path = ("params", "decoder", "moe_layers", "DeepSeekMoeBlock_0", + "MoeBlock_0", "gate", "bias") # Updates the shape to be aligned with state. moe_bias_updates = jnp.array(moe_bias_updates[0]).transpose() - new_state = maxtext_utils.update_state_param(new_state, target_path, moe_bias_updates) + new_state = maxtext_utils.update_state_param(new_state, target_path, + moe_bias_updates) else: if config.gradient_clipping_threshold > 0: - grads = maxtext_utils.apply_gradient_clipping(raw_grads, None, config.gradient_clipping_threshold) + grads = maxtext_utils.apply_gradient_clipping( + raw_grads, None, config.gradient_clipping_threshold) else: grads = raw_grads if config.optimizer_memory_host_offload: @@ -540,7 +617,8 @@ def move(path, value): # Apply updates for Auxiliary-Loss-Free load balancing for DeepSeek family if config.routed_bias and config.routed_bias_update_rate > 0.0 and moe_bias_updates is not None: target_bias = new_state.model.decoder.moe_layers.DeepSeekMoeBlock_0.MoeBlock_0.gate.bias - target_bias.value = target_bias.value + jnp.array(moe_bias_updates[0]).transpose() + target_bias.value = target_bias.value + jnp.array( + moe_bias_updates[0]).transpose() lm_loss = xent_sum / (total_weights + EPS) scalar_metrics = { @@ -555,32 +633,41 @@ def move(path, value): } if config.use_qk_clip: if isinstance(model, nn.Module): - new_state = qk_clip_utils.apply_qk_clip(new_state, intermediate_outputs, config) + new_state = qk_clip_utils.apply_qk_clip(new_state, intermediate_outputs, + config) else: - new_state = qk_clip_utils.apply_qk_clip_nnx(new_state, intermediate_outputs, config) + new_state = qk_clip_utils.apply_qk_clip_nnx(new_state, + intermediate_outputs, config) - global_max_logit = qk_clip_utils.calculate_max_logit_metric(intermediate_outputs) + global_max_logit = qk_clip_utils.calculate_max_logit_metric( + intermediate_outputs) if global_max_logit is not None: scalar_metrics["learning/max_logits"] = global_max_logit if not config.optimizer_memory_host_offload: scalar_metrics["learning/grad_norm"] = max_utils.l2norm_pytree(grads) - scalar_metrics["learning/raw_grad_norm"] = max_utils.l2norm_pytree(raw_grads) + scalar_metrics["learning/raw_grad_norm"] = max_utils.l2norm_pytree( + raw_grads) if isinstance(model, nn.Module): - scalar_metrics["learning/param_norm"] = max_utils.l2norm_pytree(new_state.params) + scalar_metrics["learning/param_norm"] = max_utils.l2norm_pytree( + new_state.params) else: model_params = nnx.state(new_state.model, nnx.Param) - scalar_metrics["learning/param_norm"] = max_utils.l2norm_pytree(model_params) + scalar_metrics["learning/param_norm"] = max_utils.l2norm_pytree( + model_params) # Surface skip-step rejections as a TB metric. The skip-step optimizer stores # is_skipped in its opt_state: the Linen path gets it from the tx.update return, # the NNX path reads it back off the optimizer it just updated in place. if config.skip_step_on_spikes: if isinstance(model, nn.Module): - is_skipped = new_opt_state.get("is_skipped") if isinstance(new_opt_state, dict) else None + is_skipped = new_opt_state.get("is_skipped") if isinstance( + new_opt_state, dict) else None else: - opt_state = nnx.to_pure_dict(nnx.state(new_state.optimizer)).get("opt_state", {}) - is_skipped = opt_state.get("is_skipped") if isinstance(opt_state, dict) else None + opt_state = nnx.to_pure_dict(nnx.state(new_state.optimizer)).get( + "opt_state", {}) + is_skipped = opt_state.get("is_skipped") if isinstance(opt_state, + dict) else None if is_skipped is not None: scalar_metrics["optim/step_skipped"] = is_skipped.astype(jnp.float32) metrics = { @@ -605,7 +692,12 @@ def eval_step(model, config, state, data, dropout_rng=None): pure_params = state.params["params"] if sparsity_enabled else state.params batch_stats = state.params.get("batch_stats", {}) - eval_loss_fn = functools.partial(loss_fn, model, config, data, dropout_rng, is_train=False) + eval_loss_fn = functools.partial(loss_fn, + model, + config, + data, + dropout_rng, + is_train=False) loss, aux = eval_loss_fn(pure_params, sparsity_state=batch_stats) else: state = nnx.merge(model, state) # reconstruct TrainStateNNX @@ -613,7 +705,8 @@ def eval_step(model, config, state, data, dropout_rng=None): mtp_acceptance_rate = 0.0 if config.mtp_eval_target_module > 0: - mtp_acceptance_rate = calculate_mtp_acceptance_rate(aux["intermediate_outputs"], config) + mtp_acceptance_rate = calculate_mtp_acceptance_rate( + aux["intermediate_outputs"], config) xent_sum = aux["xent_sum"] z_loss = aux.get("z_loss", 0.0) @@ -698,15 +791,18 @@ def training_loop_iteration( else: step_rng_args = () with maybe_record_goodput(recorder, GoodputEvent.STEP, step): - with jax.set_mesh(mesh), nn_partitioning.axis_rules(logical_axis_rules_for_train): + with jax.set_mesh(mesh), nn_partitioning.axis_rules( + logical_axis_rules_for_train): if shard_optimizer_over_data and isinstance(model, nn.Module): - state = sharding.maybe_shard_with_name(state, state_mesh_shardings, shard_mode) + state = sharding.maybe_shard_with_name(state, state_mesh_shardings, + shard_mode) state, metrics = p_train_step(state, example_batch, *step_rng_args) step_time_delta = datetime.datetime.now() - last_step_completion last_step_completion = datetime.datetime.now() - checkpointing.maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step) + checkpointing.maybe_save_checkpoint(checkpoint_manager, state, config, + data_iterator, step) if dump_hlo and step == (dump_step if dump_step >= 0 else start_step): jax.block_until_ready(state) # Ensure compilation has finished. @@ -718,7 +814,8 @@ def training_loop_iteration( all_host_upload=dump_hlo_upload_all, ) - if eval_interval > 0 and step >= start_step and step >= eval_start_step and (step + 1) % eval_interval == 0: + if eval_interval > 0 and step >= start_step and step >= eval_start_step and ( + step + 1) % eval_interval == 0: assert eval_data_iterator # Explicitly reset the eval iterator and counters before starting the eval loop eval_data_iterator.reset() @@ -731,17 +828,21 @@ def training_loop_iteration( for eval_batch in eval_data_iterator: # Shard input eval data eval_batch = jax.device_put( - eval_batch, sharding.get_input_data_sharding(config, mesh, rules=config.logical_axis_rules_for_eval) - ) + eval_batch, + sharding.get_input_data_sharding( + config, mesh, rules=config.logical_axis_rules_for_eval)) if 0 < eval_steps <= eval_step_count: break - with jax.set_mesh(mesh), nn_partitioning.axis_rules(logical_axis_rules_for_eval): + with jax.set_mesh(mesh), nn_partitioning.axis_rules( + logical_axis_rules_for_eval): eval_metrics = p_eval_step(state, eval_batch, *step_rng_args) eval_step_time_delta = datetime.datetime.now() - last_eval_step_completion last_eval_step_completion = datetime.datetime.now() metric_logger_instance.buffer_and_write_metrics( - eval_metrics, eval_step_count, step_time_delta=eval_step_time_delta, is_training=False - ) + eval_metrics, + eval_step_count, + step_time_delta=eval_step_time_delta, + is_training=False) eval_step_count += 1 prof.maybe_deactivate_profiler(step, state) @@ -749,7 +850,8 @@ def training_loop_iteration( if step == start_step: max_utils.print_mem_stats("After params initialized") - metric_logger_instance.buffer_and_write_metrics(metrics, step, step_time_delta) + metric_logger_instance.buffer_and_write_metrics(metrics, step, + step_time_delta) # Pack mutated state back to dicts jax_device_state["state"] = state @@ -776,7 +878,8 @@ def train_loop(config, recorder, state=None): # The default flag value is empty, meaning no throttling is applied by default. train_utils.maybe_apply_dcn_throttling(config) - start_step = get_first_step(model, state) # this is the start_step for training + start_step = get_first_step(model, + state) # this is the start_step for training train_utils.validate_completed_steps(start_step, config.steps) if isinstance(model, nn.Module): @@ -792,7 +895,8 @@ def train_loop(config, recorder, state=None): # the Zero-1 opt overlay doesn't apply through the diloco wrapper. params_shardings = state_mesh_shardings.params else: - params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings) + params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt( + config, state_mesh_shardings) p_train_step, p_eval_step = train_utils.jit_train_and_eval_step( config, @@ -806,11 +910,14 @@ def train_loop(config, recorder, state=None): params_shardings, ) - with jax.set_mesh(mesh), mesh, nn_partitioning.axis_rules(config.logical_axis_rules): + with jax.set_mesh(mesh), mesh, nn_partitioning.axis_rules( + config.logical_axis_rules): data_sharding = sharding.get_input_data_sharding(config, mesh) - shaped_batch = maxtext_utils.get_shaped_batch(config, batch_sharding=data_sharding) + shaped_batch = maxtext_utils.get_shaped_batch(config, + batch_sharding=data_sharding) if config.shard_optimizer_over_data and isinstance(model, nn.Module): - state = sharding.maybe_shard_with_name(state, state_mesh_shardings, config.shard_mode) + state = sharding.maybe_shard_with_name(state, state_mesh_shardings, + config.shard_mode) elif config.shard_optimizer_over_data: # NNX: reshard state so params match the data-sharded in_shardings (Zero-1 layout) state = jax.device_put(state, state_mesh_shardings) @@ -821,12 +928,15 @@ def train_loop(config, recorder, state=None): lower_args = (state, shaped_batch) maxtext_utils.maybe_dump_jaxpr(config, p_train_step, lower_args) if config.compiled_trainstep_file == "": # compile only when there is no pre-compiled file loaded - compiler_options = max_utils.parse_libtpu_flags_to_dict(config.compile_xla_flags) - compiled = p_train_step.lower(*lower_args).compile(compiler_options=compiler_options) + compiler_options = max_utils.parse_libtpu_flags_to_dict( + config.compile_xla_flags) + compiled = p_train_step.lower(*lower_args).compile( + compiler_options=compiler_options) compiled_stats = compiled.memory_analysis() max_utils.print_compiled_memory_stats(compiled_stats) prof = profiler.Profiler(config, offset_step=start_step) - metric_logger_instance = metric_logger.MetricLogger(config=config, learning_rate_schedule=learning_rate_schedule) + metric_logger_instance = metric_logger.MetricLogger( + config=config, learning_rate_schedule=learning_rate_schedule) # Write train config params, num model params, and XLA flags to tensorboard if isinstance(model, nn.Module): @@ -897,7 +1007,8 @@ def train_loop(config, recorder, state=None): state = jax_device_state["state"] if immutable_data["save_checkpoint_on_completion"]: - checkpointing.maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator) + checkpointing.maybe_save_checkpoint(checkpoint_manager, state, config, + data_iterator) if checkpoint_manager is not None: # in case the last checkpoint_period checkpoint is still in progress @@ -923,20 +1034,23 @@ def initialize(argv: Sequence[str]) -> tuple[pyconfig.HyperParameters, Any]: # TF allocates extraneous GPU memory when using TFDS data # this leads to CUDA OOMs. WAR for now is to hide GPUs from TF tf.config.set_visible_devices([], "GPU") - if "xla_tpu_spmd_rng_bit_generator_unsafe" not in os.environ.get("LIBTPU_INIT_ARGS", ""): + if "xla_tpu_spmd_rng_bit_generator_unsafe" not in os.environ.get( + "LIBTPU_INIT_ARGS", ""): os.environ["LIBTPU_INIT_ARGS"] = ( - os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true" - ) + os.environ.get("LIBTPU_INIT_ARGS", "") + + " --xla_tpu_spmd_rng_bit_generator_unsafe=true") # TODO: mazumdera@ : ensure missing mandatory fields in base.yml are filled in in argv, # or fill in here config = pyconfig.initialize(argv) max_utils.print_system_information() train_utils.validate_train_config(config) jax.config.update("jax_use_shardy_partitioner", config.shardy) - jax.config.update("jax_remove_size_one_mesh_axis_from_type", config.remove_size_one_mesh_axis_from_type) + jax.config.update("jax_remove_size_one_mesh_axis_from_type", + config.remove_size_one_mesh_axis_from_type) os.environ["TFDS_DATA_DIR"] = config.dataset_path or "" vertex_tensorboard_manager = VertexTensorboardManager() - if config.use_vertex_tensorboard or os.environ.get("UPLOAD_DATA_TO_TENSORBOARD"): + if config.use_vertex_tensorboard or os.environ.get( + "UPLOAD_DATA_TO_TENSORBOARD"): vertex_tensorboard_manager.configure_vertex_tensorboard(config) if config.use_te_comm_gemm_overlap: diff --git a/src/maxtext/utils/gradient_accumulation.py b/src/maxtext/utils/gradient_accumulation.py index 35fdf65503..579fe01049 100644 --- a/src/maxtext/utils/gradient_accumulation.py +++ b/src/maxtext/utils/gradient_accumulation.py @@ -11,7 +11,6 @@ # 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. - """Functions for gradient accumulation (GA)""" import jax @@ -34,59 +33,75 @@ def gradient_accumulation_loss_and_grad( dropout_rng, ): """ - Calculates gradients using gradient accumulation. - - This function computes the gradient of `_loss_fn` over multiple microbatches - and accumulates them before returning a single, averaged gradient. It uses - `jax.lax.scan` for efficient accumulation on device. - - It also supports a `shard_optimizer_over_data` mode (e.g., ZeRO-1) where - parameters are cast to bf16 and sharded *before* the accumulation loop - to perform the all-gather in lower precision. - - Args: - _loss_fn: The loss function to differentiate. Its signature is expected - to be: `(model, config, data, dropout_rng, params, *extra_args, is_train=True)`. - config: Model and training configuration object. Must contain - `gradient_accumulation_steps` and `shard_optimizer_over_data`. - model: The model module. - params: The model parameters (PyTree). This is only used for Linen. For NNX, - we can get the params from the model. - params_shardings: The sharding constraints for the parameters (PyTree). - data: A PyTree of batched data. The leading dimension is assumed - to be the total batch size (microbatch_size * num_accumulations). - dropout_rng: JAX PRNGKey for dropout. - - Returns: - A tuple containing: - - total_loss (Array): The mean loss, averaged over all microbatches. - - final_aux (PyTree): Auxiliary outputs, summed across microbatches. - - raw_grads (PyTree): The accumulated and averaged gradients. - """ + Calculates gradients using gradient accumulation. + + This function computes the gradient of `_loss_fn` over multiple microbatches + and accumulates them before returning a single, averaged gradient. It uses + `jax.lax.scan` for efficient accumulation on device. + + It also supports a `shard_optimizer_over_data` mode (e.g., ZeRO-1) where + parameters are cast to bf16 and sharded *before* the accumulation loop + to perform the all-gather in lower precision. + + Args: + _loss_fn: The loss function to differentiate. Its signature is expected + to be: `(model, config, data, dropout_rng, params, *extra_args, is_train=True)`. + config: Model and training configuration object. Must contain + `gradient_accumulation_steps` and `shard_optimizer_over_data`. + model: The model module. + params: The model parameters (PyTree). This is only used for Linen. For NNX, + we can get the params from the model. + params_shardings: The sharding constraints for the parameters (PyTree). + data: A PyTree of batched data. The leading dimension is assumed + to be the total batch size (microbatch_size * num_accumulations). + dropout_rng: JAX PRNGKey for dropout. + + Returns: + A tuple containing: + - total_loss (Array): The mean loss, averaged over all microbatches. + - final_aux (PyTree): Auxiliary outputs, summed across microbatches. + - raw_grads (PyTree): The accumulated and averaged gradients. + """ def _maybe_shard_with_name(inputs, sharding_names): """Wrapper of maybe_shard_with_name with fixed shard_mode""" - return maybe_shard_with_name(inputs, sharding_names, config.shard_mode, debug_sharding=config.debug_sharding) + val = inputs.get_value() if isinstance(inputs, nnx.Variable) else inputs + sharded_val = maybe_shard_with_name(val, + sharding_names, + config.shard_mode, + debug_sharding=config.debug_sharding) + return (inputs.replace( + value=sharded_val) if isinstance(inputs, nnx.Variable) else sharded_val) is_nnx = isinstance(model, nnx.Module) + _is_leaf = ((lambda x: isinstance(x, (nnx.Variable, NamedSharding))) + if is_nnx else None) + + if is_nnx: + graphdef, params, rest = nnx.split(model, nnx.Param, ...) + params = params.to_pure_dict() if hasattr(params, + "to_pure_dict") else params + params_shardings = (params_shardings.to_pure_dict() if hasattr( + params_shardings, "to_pure_dict") else params_shardings) + + def _filter_tree_to_match(target_tree, source_tree): + if isinstance(target_tree, dict) and isinstance(source_tree, dict): + return { + k: _filter_tree_to_match(target_tree[k], source_tree[k]) + for k in target_tree + if k in source_tree + } + return source_tree + + params_shardings = _filter_tree_to_match(params, params_shardings) # For ZeRO-1 + GA, read the resolved "data" axis size from the mesh rather than # config.ici_data_parallelism, which may be -1 (auto-fill) and resolves to 1 when # FSDP already consumes every device — in which case data parallelism is not active. param_mesh = jax.tree.leaves(params_shardings)[0].mesh - data_parallel_active = config.shard_mode == ShardMode.EXPLICIT and param_mesh.shape.get("data", 1) > 1 - if data_parallel_active: - # reduced/unreduced PartitionSpecs are rejected inside a jax.lax.scan carry: scan - # traces its body against an AbstractMesh whose axis types are all Auto, and the - # annotations require Explicit axes. Keep plain params_shardings in the carry and - # apply the data-parallel all-reduce to the gradients after the scan instead. - ga_params_shardings = params_shardings - grad_shardings = params_shardings - else: - ga_params_shardings = grad_shardings = params_shardings - - if is_nnx: - graphdef, params, rest = nnx.split(model, nnx.Param, ...) + data_parallel_active = (config.shard_mode == ShardMode.EXPLICIT and + param_mesh.shape.get("data", 1) > 1) + ga_params_shardings = grad_shardings = params_shardings # When using Zero-1 optimizer sharding, cast params to lower precision and apply sharding constraints # so that all-gather is done once in the lower precision before the gradient accumulation loop @@ -101,7 +116,10 @@ def convert_to_bf16(param): else: ga_params = params - ga_params = jax.tree.map(_maybe_shard_with_name, ga_params, ga_params_shardings) + ga_params = jax.tree.map(_maybe_shard_with_name, + ga_params, + ga_params_shardings, + is_leaf=_is_leaf) if is_nnx: grad_func = nnx.value_and_grad(_loss_fn, argnums=0, has_aux=True) else: @@ -116,37 +134,57 @@ def accumulate_gradient(acc_grad_and_loss, data): if is_nnx: # Reconstruct the model using the fixed parameters (ga_params) # and the advancing non-parameter state (RNGs) from the carry. - local_model = nnx.merge(graphdef, ga_params, acc_grad_and_loss["rest_state"], copy=True) + local_model = nnx.merge(graphdef, + ga_params, + acc_grad_and_loss["rest_state"], + copy=True) with set_xla_metadata(_xla_loop_unroll_strategy="double-buffer"): - (_, aux), cur_batch_gradient = grad_func(local_model, config, data, None, None, is_train=True) + (_, aux), cur_batch_gradient = grad_func(local_model, + config, + data, + None, + None, + is_train=True) + cur_batch_gradient = (cur_batch_gradient.to_pure_dict() if hasattr( + cur_batch_gradient, "to_pure_dict") else cur_batch_gradient) _, _, next_rest_state = nnx.split(local_model, nnx.Param, ...) acc_grad_and_loss["rest_state"] = next_rest_state else: - rng = ( - jax.random.fold_in(dropout_rng, acc_grad_and_loss["total_weights"].astype(jnp.int32)) - if dropout_rng is not None - else None - ) + rng = (jax.random.fold_in( + dropout_rng, acc_grad_and_loss["total_weights"].astype(jnp.int32)) + if dropout_rng is not None else None) with set_xla_metadata(_xla_loop_unroll_strategy="double-buffer"): - (_, aux), cur_batch_gradient = grad_func(model, config, data, rng, ga_params, is_train=True) + (_, aux), cur_batch_gradient = grad_func(model, + config, + data, + rng, + ga_params, + is_train=True) acc_grad_and_loss["loss"] += aux["xent_sum"] + aux.get("dpo_loss", 0.0) acc_grad_and_loss["moe_lb_loss"] += aux["moe_lb_loss"] acc_grad_and_loss["indexer_loss"] += aux["indexer_loss"] acc_grad_and_loss["mtp_loss"] += aux["mtp_loss"] - acc_grad_and_loss["grad"] = jax.tree_util.tree_map(lambda x, y: x + y, cur_batch_gradient, acc_grad_and_loss["grad"]) + acc_grad_and_loss["grad"] = jax.tree_util.tree_map( + lambda x, y: x + y, cur_batch_gradient, acc_grad_and_loss["grad"]) acc_grad_and_loss["total_weights"] += aux["total_weights"] return acc_grad_and_loss, aux def reshape_to_microbatch_accumulations(batch_arr): """Reshape global batch to microbatches, assuming batch axis is leading.""" num_microbatches = config.gradient_accumulation_steps - microbatch_shape = (batch_arr.shape[0] // num_microbatches, num_microbatches) + batch_arr.shape[1:] + microbatch_shape = ( + batch_arr.shape[0] // num_microbatches, + num_microbatches, + ) + batch_arr.shape[1:] reshaped_batch_arr = jnp.reshape(batch_arr, microbatch_shape) return jnp.swapaxes(reshaped_batch_arr, 0, 1) data = jax.tree_util.tree_map(reshape_to_microbatch_accumulations, data) init_grad = jax.tree_util.tree_map(jnp.zeros_like, ga_params) - init_grad = jax.tree.map(_maybe_shard_with_name, init_grad, grad_shardings) + init_grad = jax.tree.map(_maybe_shard_with_name, + init_grad, + grad_shardings, + is_leaf=_is_leaf) init_grad_and_loss = { "loss": 0.0, # accumulates xent_sum across microbatches "grad": init_grad, @@ -160,27 +198,39 @@ def reshape_to_microbatch_accumulations(batch_arr): init_grad_and_loss["rest_state"] = rest # pyrefly: ignore[unbound-name] grad_and_loss, aux = jax.lax.scan( - accumulate_gradient, init_grad_and_loss, data, length=config.gradient_accumulation_steps - ) - loss = ( - grad_and_loss["loss"] / grad_and_loss["total_weights"] - + grad_and_loss["moe_lb_loss"] / config.gradient_accumulation_steps - + grad_and_loss["indexer_loss"] / config.gradient_accumulation_steps - + grad_and_loss["mtp_loss"] / config.gradient_accumulation_steps + accumulate_gradient, + init_grad_and_loss, + data, + length=config.gradient_accumulation_steps, ) + loss = (grad_and_loss["loss"] / grad_and_loss["total_weights"] + + grad_and_loss["moe_lb_loss"] / config.gradient_accumulation_steps + + grad_and_loss["indexer_loss"] / config.gradient_accumulation_steps + + grad_and_loss["mtp_loss"] / config.gradient_accumulation_steps) raw_grads = grad_and_loss["grad"] if data_parallel_active: # Mark the gradients unreduced over the "data" axis now that we're outside the # scan; this triggers the cross-replica all-reduce. The annotation can't live in # the scan carry (see above), so it's applied here instead. - unreduced_shardings = jax.tree.map(update_sharding_for_unreduced, params_shardings) - raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, unreduced_shardings) - raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, params_shardings) - raw_grads = jax.tree_util.tree_map(lambda arr: arr / grad_and_loss["total_weights"], raw_grads) + unreduced_shardings = jax.tree.map(update_sharding_for_unreduced, + params_shardings) + raw_grads = jax.tree.map(_maybe_shard_with_name, + raw_grads, + unreduced_shardings, + is_leaf=_is_leaf) + raw_grads = jax.tree.map(_maybe_shard_with_name, + raw_grads, + params_shardings, + is_leaf=_is_leaf) + raw_grads = jax.tree_util.tree_map( + lambda arr: arr / grad_and_loss["total_weights"], raw_grads) aux = jax.tree.map(lambda x: jnp.sum(x, axis=0), aux) # pytype: disable=module-attr if is_nnx: nnx.update(model, grad_and_loss["rest_state"]) + raw_grads_state = nnx.state(model, nnx.Param) + raw_grads_state.replace_by_pure_dict(raw_grads) + raw_grads = raw_grads_state return loss, aux, raw_grads @@ -188,13 +238,13 @@ def reshape_to_microbatch_accumulations(batch_arr): # GA helper functions def update_sharding_for_reduced(sharding: NamedSharding) -> NamedSharding: """ - Add reduced on data axis of given NamedSharding - """ + Add reduced on data axis of given NamedSharding + """ return sharding.update(spec=sharding.spec.update(reduced={"data"})) def update_sharding_for_unreduced(sharding: NamedSharding) -> NamedSharding: """ - Add unreduced on data axis of given NamedSharding - """ + Add unreduced on data axis of given NamedSharding + """ return sharding.update(spec=sharding.spec.update(unreduced={"data"})) diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py index 05d95cd0f4..555d61a3b0 100644 --- a/src/maxtext/utils/lora_utils.py +++ b/src/maxtext/utils/lora_utils.py @@ -11,7 +11,6 @@ # 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. - """Common LoRA utils needed to support LoRA adapters.""" from collections.abc import Mapping @@ -34,7 +33,6 @@ from maxtext.utils import gcs_utils from maxtext.utils import max_logging from maxtext.utils import max_utils -from maxtext.utils import maxtext_utils from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR # NNX-only imports (train_state_nnx, model_creation_utils) are loaded lazily @@ -57,17 +55,21 @@ def apply_lora_on_base_params(base_params, lora_params, lora_scale_factor=1.0): def lora_update_or_base(base_weight, lora_a, lora_b): if lora_a is not None and lora_b is not None: - return base_weight + jnp.einsum("br,rnd->bnd", lora_b, lora_a) * lora_scale_factor + return base_weight + jnp.einsum("br,rnd->bnd", lora_b, + lora_a) * lora_scale_factor else: return base_weight # Keep the base weight if no Lora update def apply_lora_recursively(base_params, lora_params, module_name): for name, param in lora_params.items(): if isinstance(param, dict): - apply_lora_recursively(base_params[name], param, f"{module_name}.{name}") + apply_lora_recursively(base_params[name], param, + f"{module_name}.{name}") elif param is not None: if name not in ["lora_a.kernel", "lora_b.kernel"]: - raise ValueError(f"Unexpected non-lora specific weights ({module_name}.{name}) found in the lora_params") + raise ValueError( + f"Unexpected non-lora specific weights ({module_name}.{name}) found in the lora_params" + ) lora_b = lora_params["lora_a.kernel"] lora_a = lora_params["lora_b.kernel"] @@ -80,7 +82,9 @@ def apply_lora_recursively(base_params, lora_params, module_name): apply_lora_recursively(base_params, lora_params, "") -def unapply_lora_from_base_params(base_params, lora_params, lora_scale_factor=1.0): +def unapply_lora_from_base_params(base_params, + lora_params, + lora_scale_factor=1.0): """ Unapply the LoRA weights from the base weights of the model using formula: W_org = W - BA, where @@ -96,17 +100,21 @@ def unapply_lora_from_base_params(base_params, lora_params, lora_scale_factor=1. def lora_update_or_base(base_weight, lora_a, lora_b): if lora_a is not None and lora_b is not None: - return base_weight - jnp.einsum("br,rnd->bnd", lora_b, lora_a) * lora_scale_factor + return base_weight - jnp.einsum("br,rnd->bnd", lora_b, + lora_a) * lora_scale_factor else: return base_weight # Keep the base weight if no Lora update def unapply_lora_recursively(base_params, lora_params, module_name): for name, param in lora_params.items(): if isinstance(param, dict): - unapply_lora_recursively(base_params[name], param, f"{module_name}.{name}") + unapply_lora_recursively(base_params[name], param, + f"{module_name}.{name}") elif param is not None: if name not in ["lora_a.kernel", "lora_b.kernel"]: - raise ValueError(f"Unexpected non-lora specific weights ({module_name}.{name}) found in the lora_params") + raise ValueError( + f"Unexpected non-lora specific weights ({module_name}.{name}) found in the lora_params" + ) lora_b = lora_params["lora_a.kernel"] lora_a = lora_params["lora_b.kernel"] @@ -119,7 +127,8 @@ def unapply_lora_recursively(base_params, lora_params, module_name): unapply_lora_recursively(base_params, lora_params, "") -def load_adapter(config, base_abstract_state_params, adapter_config_path, adapter_weights_path): +def load_adapter(config, base_abstract_state_params, adapter_config_path, + adapter_weights_path): """Load a LoRA adapter from disk and return its parameters. When `config.pure_nnx` is True, `base_abstract_state_params` is the NNX @@ -147,15 +156,20 @@ def load_adapter(config, base_abstract_state_params, adapter_config_path, adapte lora_config = json.load(f) if lora_config is None: - raise FileNotFoundError(f"Failed to read lora_config from {adapter_config_path}.") + raise FileNotFoundError( + f"Failed to read lora_config from {adapter_config_path}.") - if not gcs_utils.gcs_path_exists(f"{adapter_weights_path}/commit_success.txt"): - raise FileNotFoundError(f"Failed to read lora_weights from {adapter_weights_path}.") + if not gcs_utils.gcs_path_exists( + f"{adapter_weights_path}/commit_success.txt"): + raise FileNotFoundError( + f"Failed to read lora_weights from {adapter_weights_path}.") if config.pure_nnx: - lora_state, _ = get_lora_abstract_state_nnx(base_abstract_state_params, lora_config) + lora_state, _ = get_lora_abstract_state_nnx(base_abstract_state_params, + lora_config) else: - lora_state, _ = get_lora_abstract_state(base_abstract_state_params, lora_config) + lora_state, _ = get_lora_abstract_state(base_abstract_state_params, + lora_config) with nn_partitioning.axis_rules(config.logical_axis_rules): lora_params = checkpointing.load_params_from_path( @@ -169,7 +183,8 @@ def load_adapter(config, base_abstract_state_params, adapter_config_path, adapte return lora_params, lora_config -def setup_initial_lora_state(model, data_iterator, tx, config, rng, mesh, checkpoint_manager, lora_adapter_path): +def setup_initial_lora_state(model, data_iterator, tx, config, rng, mesh, + checkpoint_manager, lora_adapter_path): """Initialize the LoRA train state and optionally restore it from a checkpoint. On the NNX path, `model` is unused; the abstract state is built from @@ -198,13 +213,16 @@ def setup_initial_lora_state(model, data_iterator, tx, config, rng, mesh, checkp lora_config = None if lora_adapter_path: - max_logging.log(f"Setting initial state of LoRA with lora_adapter_path = {lora_adapter_path}") + max_logging.log( + f"Setting initial state of LoRA with lora_adapter_path = {lora_adapter_path}" + ) if config.pure_nnx: # pylint: disable=import-outside-toplevel from maxtext.common import train_state_nnx from maxtext.utils import model_creation_utils - _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh) + _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model( + config, mesh) def create_train_state_fn(): nnx_model = _create_model_partial() @@ -213,8 +231,14 @@ def create_train_state_fn(): init_state_fn = create_train_state_fn else: - init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, True, rng) - unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(config, mesh, init_state_fn, True) + from maxtext.utils import maxtext_utils # pylint: disable=import-outside-toplevel + + init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, + config, True, rng) + from maxtext.utils import maxtext_utils # pylint: disable=import-outside-toplevel + + unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state( + config, mesh, init_state_fn, True) lora_config_path = lora_adapter_path + "adapter_config.json" @@ -222,9 +246,11 @@ def create_train_state_fn(): if config.pure_nnx: base_abstract_params = _nnx_param_subtree(unboxed_abstract_state) - lora_state, lora_state_annotations = get_lora_abstract_state_nnx(base_abstract_params, lora_config) + lora_state, lora_state_annotations = get_lora_abstract_state_nnx( + base_abstract_params, lora_config) else: - lora_state, lora_state_annotations = get_lora_abstract_state(unboxed_abstract_state.params, lora_config) + lora_state, lora_state_annotations = get_lora_abstract_state( + unboxed_abstract_state.params, lora_config) lora_weights_path = f"{lora_adapter_path}/0/items" @@ -243,7 +269,8 @@ def create_train_state_fn(): ) if restored_lora: - raise NotImplementedError("This codepath is not implemented for LoRA adapters yet.") + raise NotImplementedError( + "This codepath is not implemented for LoRA adapters yet.") else: lora_state = lora_state.replace(params=raw_lora_params) lora_state = max_utils.unbox_logicallypartioned(lora_state) @@ -275,7 +302,9 @@ def get_lora_abstract_state(base_abstract_params, lora_config): } lora_target_modules = lora_config["target_modules"] - lora_target_modules = [other_lora_format_to_jax_format.get(s, s) for s in lora_target_modules] + lora_target_modules = [ + other_lora_format_to_jax_format.get(s, s) for s in lora_target_modules + ] lora_rank = int(lora_config["r"]) @@ -314,7 +343,9 @@ def get_lora_param_sharding(base_param_sharding, lora_module): base_sharding_pspec_size = len(base_param_sharding.spec) if base_sharding_pspec_size > 4: - raise ValueError("Encountered unexpected size of PartitionSpec in sharding. Size > 4 is not supported") + raise ValueError( + "Encountered unexpected size of PartitionSpec in sharding. Size > 4 is not supported" + ) base_mesh = base_param_sharding.mesh base_memory_kind = base_param_sharding.memory_kind @@ -335,14 +366,19 @@ def get_lora_param_sharding(base_param_sharding, lora_module): lora_a_pspec_tuple = base_pspec[:-1] + ((),) lora_a_pspec = jax.sharding.PartitionSpec(*lora_a_pspec_tuple) if base_sharding_pspec_size == 4: - lora_b_pspec = jax.sharding.PartitionSpec((), base_pspec[1], base_pspec[-1]) + lora_b_pspec = jax.sharding.PartitionSpec((), base_pspec[1], + base_pspec[-1]) else: lora_b_pspec = jax.sharding.PartitionSpec((), base_pspec[-1]) else: raise ValueError(f"Unsupported lora_module={lora_module}") - lora_a_sharding = jax.sharding.NamedSharding(mesh=base_mesh, spec=lora_a_pspec, memory_kind=base_memory_kind) - lora_b_sharding = jax.sharding.NamedSharding(mesh=base_mesh, spec=lora_b_pspec, memory_kind=base_memory_kind) + lora_a_sharding = jax.sharding.NamedSharding(mesh=base_mesh, + spec=lora_a_pspec, + memory_kind=base_memory_kind) + lora_b_sharding = jax.sharding.NamedSharding(mesh=base_mesh, + spec=lora_b_pspec, + memory_kind=base_memory_kind) return lora_a_sharding, lora_b_sharding @@ -361,7 +397,8 @@ def module_is_target_module(module, target_modules): return target_module return None - def add_lora_params(lora_params, module_name, base_params, lora_rank, lora_target_modules): + def add_lora_params(lora_params, module_name, base_params, lora_rank, + lora_target_modules): for name, param in base_params.items(): if isinstance(param, dict): lora_params[name] = {} @@ -374,38 +411,59 @@ def add_lora_params(lora_params, module_name, base_params, lora_rank, lora_targe ) else: if name not in ["kernel", "scale", "embedding"]: - raise ValueError(f"Unexpected key={name} exists in the abstract params of base model.") + raise ValueError( + f"Unexpected key={name} exists in the abstract params of base model." + ) if not isinstance(param, jax.ShapeDtypeStruct): - raise ValueError("Unexpected type found in the abstract params of the base model.") + raise ValueError( + "Unexpected type found in the abstract params of the base model.") lora_a_key = "lora_a.kernel" lora_b_key = "lora_b.kernel" - target_module = module_is_target_module(module_name, lora_target_modules) + target_module = module_is_target_module(module_name, + lora_target_modules) if target_module is not None: - lora_a_shape, lora_b_shape = get_lora_param_shape(param.shape, lora_rank, target_module) + lora_a_shape, lora_b_shape = get_lora_param_shape( + param.shape, lora_rank, target_module) base_dtype = param.dtype - lora_a_sharding, lora_b_sharding = get_lora_param_sharding(param.sharding, target_module) + lora_a_sharding, lora_b_sharding = get_lora_param_sharding( + param.sharding, target_module) - lora_params[lora_a_key] = jax.ShapeDtypeStruct(shape=lora_a_shape, dtype=base_dtype, sharding=lora_a_sharding) + lora_params[lora_a_key] = jax.ShapeDtypeStruct( + shape=lora_a_shape, dtype=base_dtype, sharding=lora_a_sharding) - lora_params[lora_b_key] = jax.ShapeDtypeStruct(shape=lora_b_shape, dtype=base_dtype, sharding=lora_b_sharding) + lora_params[lora_b_key] = jax.ShapeDtypeStruct( + shape=lora_b_shape, dtype=base_dtype, sharding=lora_b_sharding) else: lora_params[name] = None def get_lora_annotations(lora_abstract_params): - return jax.tree_util.tree_map(lambda x: x.sharding.spec, lora_abstract_params) + return jax.tree_util.tree_map( + lambda x: x.sharding.spec + if getattr(x, "sharding", None) is not None else None, + lora_abstract_params, + ) - add_lora_params(lora_abstract_params, "", base_abstract_params, lora_rank, lora_target_modules) + add_lora_params(lora_abstract_params, "", base_abstract_params, lora_rank, + lora_target_modules) unboxed_abstract_lora_state = train_state.TrainState( - step=0, apply_fn=None, params=lora_abstract_params, tx=None, opt_state={} # type: ignore + step=0, + apply_fn=None, + params=lora_abstract_params, + tx=None, + opt_state={} # type: ignore ) lora_state_mesh_annotations = train_state.TrainState( - step=0, apply_fn=None, params=get_lora_annotations(lora_abstract_params), tx=None, opt_state={} # type: ignore + step=0, + apply_fn=None, + params=get_lora_annotations(lora_abstract_params), + tx=None, + opt_state={} # type: ignore ) return unboxed_abstract_lora_state, lora_state_mesh_annotations @@ -419,30 +477,57 @@ def _get_lora_module_path(mt_config: pyconfig.HyperParameters) -> str: if mt_config.lora.lora_module_path: return mt_config.lora.lora_module_path - config_path = os.path.join(MAXTEXT_CONFIGS_DIR, "post_train", "lora_module_path.yml") + config_path = os.path.join(MAXTEXT_CONFIGS_DIR, "post_train", + "lora_module_path.yml") lora_configs = pyconfig._load_config(config_path) # pylint: disable=protected-access model_name = mt_config.model_name.lower() - - # Find the first matching architecture prefix or use 'default' - matched_key = next((k for k in lora_configs if k != "default" and model_name.startswith(k)), "default") + # Find matching model name or architecture prefix in lora_module_path.yml + matched_key = "default" + for key in lora_configs: + if key == "default": + continue + key_lower = key.lower() + if model_name == key_lower or model_name.startswith(key_lower): + matched_key = key + break + + # Fallback to normalized match (ignoring hyphens and underscores) if direct prefix match fails + if matched_key == "default": + model_name_norm = model_name.replace("-", "").replace("_", "") + for key in lora_configs: + if key == "default": + continue + key_norm = key.lower().replace("-", "").replace("_", "") + if model_name_norm == key_norm or model_name_norm.startswith(key_norm): + matched_key = key + break if matched_key == "default": - max_logging.log(f"Warning: Model '{model_name}' is unverified; falling back to default LoRA path.") + max_logging.log( + f"Warning: Model '{model_name}' is unverified; falling back to default LoRA path." + ) else: - max_logging.log(f"Auto-detected lora_module_path for model '{model_name}' (matched: '{matched_key}')") + max_logging.log( + f"Auto-detected lora_module_path for model '{model_name}' (matched: '{matched_key}')" + ) - raw_path = lora_configs.get(matched_key, "decoder/layers/.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))") + raw_path = lora_configs.get( + matched_key, + "decoder/layers/.*(self_attention/(query|key|value|out)|mlp/(wi_0|wi_1|wo))" + ) # This regex makes the layer index optional, matching scanned, unscanned named (layers_0), - # and unscanned index (layers/0) layer paths. - layer_pattern = r"layers(?:_[0-9]+|/[0-9]+)?/" - final_path = str(raw_path).replace("layers/", layer_pattern) + # unscanned index (layers/0), and alternative block prefixes (blocks, decoder_layers). + layer_pattern = r"(?:_[0-9]+|/[0-9]+)?/" + final_path = re.sub(r"(layers|blocks|decoder_layers)/", r"\1" + layer_pattern, + str(raw_path)) max_logging.log(f"Using lora_module_path: {final_path}") return final_path -def _build_lora_provider(mt_config: pyconfig.HyperParameters) -> qwix.LoraProvider: +def _build_lora_provider( + mt_config: pyconfig.HyperParameters) -> qwix.LoraProvider: """Builds a Qwix LoRA provider from MaxText LoRA settings.""" lora_module_path = _get_lora_module_path(mt_config) lora_kwargs = { @@ -482,7 +567,8 @@ def is_lora_enabled(model: nnx.Module) -> bool: return False -def _verify_lora_parameters(lora_model: nnx.Module, mt_config: pyconfig.HyperParameters) -> None: +def _verify_lora_parameters(lora_model: nnx.Module, + mt_config: pyconfig.HyperParameters) -> None: """Validates that LoRA is active or that target modules were matched.""" enabled = is_lora_enabled(lora_model) @@ -498,13 +584,16 @@ def _verify_lora_parameters(lora_model: nnx.Module, mt_config: pyconfig.HyperPar wrapped_modules = sorted(list(wrapped_modules)) max_logging.log( f"LoRA configured: module_path='{_get_lora_module_path(mt_config)}' successfully matched " - f"{len(wrapped_modules)} target submodules." - ) + f"{len(wrapped_modules)} target submodules.") preview_limit = 20 preview_modules = wrapped_modules[:preview_limit] - max_logging.log(f"Sample matched submodules ({len(preview_modules)} of {len(wrapped_modules)}): {preview_modules}") + max_logging.log( + f"Sample matched submodules ({len(preview_modules)} of {len(wrapped_modules)}): {preview_modules}" + ) else: - max_logging.log("LoRA is enabled. (Detailed submodules match report skipped due to mock model or empty state)") + max_logging.log( + "LoRA is enabled. (Detailed submodules match report skipped due to mock model or empty state)" + ) return lora_module_path = _get_lora_module_path(mt_config) @@ -517,22 +606,28 @@ def _verify_lora_parameters(lora_model: nnx.Module, mt_config: pyconfig.HyperPar matched_module_paths.append(module_path) if not matched_module_paths: - max_logging.log(f"Error: LoRA module_path='{lora_module_path}' did not match any weights.") - raise ValueError("LoRA enabled but no LoRA parameters found in decoder/model state.") + max_logging.log( + f"Error: LoRA module_path='{lora_module_path}' did not match any weights." + ) + raise ValueError( + "LoRA enabled but no LoRA parameters found in decoder/model state.") # Simplify matched paths by replacing numeric layer indices with "*" to avoid redundant output - simplified_matches = sorted( - {"/".join("*" if p.isdigit() else p for p in path.split("/")) for path in matched_module_paths} + simplified_matches = sorted({ + "/".join("*" if p.isdigit() else p + for p in path.split("/")) + for path in matched_module_paths + }) + max_logging.log( + f"LoRA target verification: successfully matched {len(matched_module_paths)} modules." ) - max_logging.log(f"LoRA target verification: successfully matched {len(matched_module_paths)} modules.") max_logging.log(f"Matched submodule patterns: {simplified_matches}") raise ValueError( "LoRA module path matched target modules, but nnx.LoRAParam is still " "missing. For Tunix PeftTrainer, LoRA params must be materialized before " "trainer initialization, otherwise it falls back to full-model training. " - f"Sample matches: {matched_module_paths[:10]}" - ) + f"Sample matches: {matched_module_paths[:10]}") def sync_lora_metadata(config: pyconfig.HyperParameters) -> None: @@ -570,8 +665,7 @@ def sync_lora_metadata(config: pyconfig.HyperParameters) -> None: config.lora.lora_alpha = meta_alpha max_logging.log( f"Synced LoRA parameters from Orbax metadata at {lora_restore_path}: " - f"rank={config.lora.lora_rank}, alpha={config.lora.lora_alpha}" - ) + f"rank={config.lora.lora_rank}, alpha={config.lora.lora_alpha}") def apply_lora_to_model( @@ -583,7 +677,8 @@ def apply_lora_to_model( # pylint: disable=protected-access # Skip Qwix LoRA if MaxText LoRA adapters are loaded if mt_config.lora_input_adapters_path: - max_logging.log("MaxText LoRA adapters loaded, skipping Qwix LoRA application") + max_logging.log( + "MaxText LoRA adapters loaded, skipping Qwix LoRA application") return model if not mt_config.lora.enable_lora: @@ -593,12 +688,16 @@ def apply_lora_to_model( lora_provider = _build_lora_provider(mt_config) - dp_size = 1 - if mesh is not None and "data" in mesh.shape: - dp_size = mesh.shape["data"] - model_rngs = getattr(model.decoder, "rngs", None) - decoder_input_tokens, decoder_positions = _prepare_dummy_inputs(dummy_bs=dp_size) + batch_partition_size = 1 + if mesh is not None and hasattr(mesh, "shape"): + batch_partition_size = mesh.shape.get("data", 1) * mesh.shape.get("fsdp", 1) + dummy_bs = int( + max( + getattr(mt_config, "per_device_batch_size", 1) * batch_partition_size, + batch_partition_size)) + decoder_input_tokens, decoder_positions = _prepare_dummy_inputs( + dummy_bs=dummy_bs) lora_model = qwix.apply_lora_to_model( model, @@ -609,24 +708,32 @@ def apply_lora_to_model( ) if mesh is not None: - with jax.set_mesh(mesh), nn_partitioning.axis_rules(mt_config.logical_axis_rules): + + def _apply_sharding(): + nonlocal lora_model graph_def, state = nnx.split(lora_model) # We handle explicit replication for LoRA to ensure safety and efficiency. state = jax.tree_util.tree_map( - lambda x: x.replace(sharding=jax.sharding.PartitionSpec(), out_sharding=None, sharding_names=None) - if isinstance(x, nnx.LoRAParam) - else x, + lambda x: x.replace(sharding=jax.sharding.PartitionSpec(), + out_sharding=None, + sharding_names=None) + if isinstance(x, (nnx.LoRAParam, nnx.Param, nnx.Variable)) and + (isinstance(x, nnx.LoRAParam) or getattr(x, "sharding", None) is None + ) else x, state, is_leaf=lambda x: isinstance(x, nnx.Variable), ) # Use logical_to_mesh_sharding to correctly map logical axes like 'embed' # to physical mesh axes. - dst_shardings = nn.logical_to_mesh_sharding(nnx.get_partition_spec(state), mesh, mt_config.logical_axis_rules) + dst_shardings = nn.logical_to_mesh_sharding(nnx.get_partition_spec(state), + mesh, + mt_config.logical_axis_rules) def _safe_reshard(var, sharding_spec): - if not isinstance(var, nnx.Variable) or not isinstance(sharding_spec, jax.sharding.Sharding): + if not isinstance(var, nnx.Variable) or not isinstance( + sharding_spec, jax.sharding.Sharding): return var val = var.get_value() if not isinstance(val, jax.Array): @@ -634,19 +741,72 @@ def _safe_reshard(var, sharding_spec): # 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. - resharded_val = jax.make_array_from_callback(val.shape, sharding_spec, lambda idx: val[idx]) + resharded_val = jax.make_array_from_callback(val.shape, sharding_spec, + lambda idx: val[idx]) return var.replace(value=resharded_val) - state = jax.tree_util.tree_map(_safe_reshard, state, dst_shardings, is_leaf=lambda x: isinstance(x, nnx.Variable)) + state = jax.tree_util.tree_map( + _safe_reshard, + state, + dst_shardings, + is_leaf=lambda x: isinstance(x, nnx.Variable)) lora_model = nnx.merge(graph_def, state) - _verify_lora_parameters(lora_model, mt_config) + try: + with jax.set_mesh(mesh), nn_partitioning.axis_rules( + mt_config.logical_axis_rules): + _apply_sharding() + except ValueError: + with nn_partitioning.axis_rules(mt_config.logical_axis_rules): + _apply_sharding() + _verify_lora_parameters(lora_model, mt_config) + nnx.pop(lora_model, nnx.Intermediate) return lora_model -def restore_lora_from_path(model: nnx.Module, mt_config: pyconfig.HyperParameters) -> nnx.Module: +def reshard_obj_to_mesh(obj, mesh, mt_config): + """Reshards an NNX model or optimizer object to the given mesh using its partition specs. + + Args: + obj: The NNX model or optimizer state object. + mesh: The JAX mesh object for parameter sharding. + mt_config: The HyperParameters configuration containing logical axis rules. + + Returns: + The resharded NNX object. + """ + if mesh is None: + return obj + graph_def, state = nnx.split(obj) + dst_shardings = nn.logical_to_mesh_sharding(nnx.get_partition_spec(state), + mesh, + mt_config.logical_axis_rules) + + def _safe_reshard(var, sharding_spec): + if not isinstance(var, nnx.Variable): + return var + spec = sharding_spec.value if hasattr(sharding_spec, + "value") else sharding_spec + if not isinstance(spec, jax.sharding.Sharding): + return var + val = var.value + if isinstance(val, jax.Array) and val.sharding != spec: + resharded_val = jax.make_array_from_callback(val.shape, spec, + lambda idx: val[idx]) + return var.replace(value=resharded_val) + return var + + state = jax.tree_util.tree_map(_safe_reshard, + state, + dst_shardings, + is_leaf=lambda x: isinstance(x, nnx.Variable)) + return nnx.merge(graph_def, state) + + +def restore_lora_from_path(model: nnx.Module, + mt_config: pyconfig.HyperParameters) -> nnx.Module: """Restores LoRA parameter weights from an external Orbax checkpoint. This function performs the restore in-place on the model's parameters and @@ -679,13 +839,20 @@ def restore_lora_from_path(model: nnx.Module, mt_config: pyconfig.HyperParameter abstract_lora_params = nnx.state(model, nnx.LoRAParam) target_for_restore = jax.tree.map( - lambda v: {"value": v.value}, + lambda v: { + "value": + v.get_value() if hasattr(v, "get_value") else v[...] + if hasattr(v, "__getitem__") else v + }, abstract_lora_params, is_leaf=lambda n: isinstance(n, nnx.Variable), ) - sharding_tree = jax.tree.map(lambda x: x.sharding if hasattr(x, "sharding") else None, target_for_restore) - restore_args_tree = ocp.checkpoint_utils.construct_restore_args(target_for_restore, sharding_tree) + sharding_tree = jax.tree.map( + lambda x: x.sharding + if hasattr(x, "sharding") else None, target_for_restore) + restore_args_tree = ocp.checkpoint_utils.construct_restore_args( + target_for_restore, sharding_tree) try: restore_args = ocp.args.PyTreeRestore( @@ -698,41 +865,11 @@ def restore_lora_from_path(model: nnx.Module, mt_config: pyconfig.HyperParameter args=restore_args, ) except Exception as e: # pylint: disable=broad-exception-caught - max_logging.log(f"Guided restore failed: {e}. Falling back to basic restore.") + max_logging.log( + f"Guided restore failed: {e}. Falling back to basic restore.") restored_lora_params = ocp.PyTreeCheckpointer().restore(lora_restore_path) - # Post processing - def _map_to_state(path, variable): - if not isinstance(variable, nnx.Variable): - return - - str_path = [str(k.key if hasattr(k, "key") else (k.name if hasattr(k, "name") else k)) for k in path] - - curr = restored_lora_params - for p in str_path: - if isinstance(curr, dict) and p in curr: - curr = curr[p] - elif hasattr(curr, p): - curr = getattr(curr, p) - else: - return - - if isinstance(curr, dict) and "value" in curr: - matched_val = curr["value"] - elif hasattr(curr, "value"): - matched_val = getattr(curr, "value") - else: - matched_val = curr - - variable.value = matched_val - - jax.tree_util.tree_map_with_path( - _map_to_state, - abstract_lora_params, - is_leaf=lambda n: isinstance(n, nnx.Variable), - ) - - nnx.update(model, abstract_lora_params) + checkpointing._update_nnx_state_from_pure_dict(model, restored_lora_params) # pylint: disable=protected-access max_logging.log(f"LoRA restore complete from '{lora_restore_path}'.") return model @@ -752,10 +889,13 @@ def _is_nnx_branch(x): def _nnx_param_subtree(unboxed_abstract_state): """Return the `model` substate, peeling off the outer `TrainStateNNX` wrapping.""" - return unboxed_abstract_state["model"] if "model" in unboxed_abstract_state else unboxed_abstract_state + return unboxed_abstract_state[ + "model"] if "model" in unboxed_abstract_state else unboxed_abstract_state -def apply_lora_on_base_params_nnx(base_params, lora_params, lora_scale_factor=1.0): +def apply_lora_on_base_params_nnx(base_params, + lora_params, + lora_scale_factor=1.0): """Apply LoRA deltas to `base_params` on the NNX path. Standard LoRA decomposition: `W_new = W + lora_a @ lora_b * scale`, where @@ -776,18 +916,22 @@ def recurse(base_node, lora_node, path): lora_a = lora_node["lora_a.kernel"] lora_b = lora_node["lora_b.kernel"] if lora_a is not None and lora_b is not None: - base_node["kernel"] = base_node["kernel"] + jnp.einsum("er,rnd->end", lora_a, lora_b) * lora_scale_factor + base_node["kernel"] = base_node["kernel"] + jnp.einsum( + "er,rnd->end", lora_a, lora_b) * lora_scale_factor return for name, lora_child in lora_node.items(): if _is_nnx_branch(lora_child): recurse(base_node[name], lora_child, f"{path}.{name}") elif lora_child is not None: - raise ValueError(f"Unexpected non-lora key ({path}.{name}) in lora_params") + raise ValueError( + f"Unexpected non-lora key ({path}.{name}) in lora_params") recurse(base_params, lora_params, "") -def unapply_lora_from_base_params_nnx(base_params, lora_params, lora_scale_factor=1.0): +def unapply_lora_from_base_params_nnx(base_params, + lora_params, + lora_scale_factor=1.0): """Unapply LoRA deltas from `base_params` on the NNX path. Symmetric inverse of `apply_lora_on_base_params_nnx`: `W -= lora_a @ lora_b * scale` @@ -800,13 +944,15 @@ def recurse(base_node, lora_node, path): lora_a = lora_node["lora_a.kernel"] lora_b = lora_node["lora_b.kernel"] if lora_a is not None and lora_b is not None: - base_node["kernel"] = base_node["kernel"] - jnp.einsum("er,rnd->end", lora_a, lora_b) * lora_scale_factor + base_node["kernel"] = base_node["kernel"] - jnp.einsum( + "er,rnd->end", lora_a, lora_b) * lora_scale_factor return for name, lora_child in lora_node.items(): if _is_nnx_branch(lora_child): recurse(base_node[name], lora_child, f"{path}.{name}") elif lora_child is not None: - raise ValueError(f"Unexpected non-lora key ({path}.{name}) in lora_params") + raise ValueError( + f"Unexpected non-lora key ({path}.{name}) in lora_params") recurse(base_params, lora_params, "") @@ -835,13 +981,17 @@ def get_lora_abstract_state_nnx(base_abstract_params, lora_config): "o_proj": "self_attention.out", } - lora_target_modules = [other_lora_format_to_jax_format.get(s, s) for s in lora_config["target_modules"]] + lora_target_modules = [ + other_lora_format_to_jax_format.get(s, s) + for s in lora_config["target_modules"] + ] lora_rank = int(lora_config["r"]) def get_lora_param_shape(base_array_shape, lora_module): if len(base_array_shape) > 4: raise ValueError(f"Unsupported base array shape {base_array_shape} (>4D)") - if lora_module in ("self_attention.query", "self_attention.key", "self_attention.value"): + if lora_module in ("self_attention.query", "self_attention.key", + "self_attention.value"): lora_a_shape = base_array_shape[:-2] + (lora_rank,) lora_b_shape = (lora_rank,) + base_array_shape[1:] elif lora_module == "self_attention.out": @@ -860,13 +1010,15 @@ def get_lora_param_sharding(base_param_sharding, lora_module): base_pspec = base_param_sharding.spec if len(base_pspec) > 4: raise ValueError("PartitionSpec size > 4 not supported") - if lora_module in ("self_attention.query", "self_attention.key", "self_attention.value"): + if lora_module in ("self_attention.query", "self_attention.key", + "self_attention.value"): lora_a_pspec = jax.sharding.PartitionSpec(*(base_pspec[:-2] + ((),))) lora_b_pspec = jax.sharding.PartitionSpec(*(((),) + base_pspec[1:])) elif lora_module == "self_attention.out": lora_a_pspec = jax.sharding.PartitionSpec(*(base_pspec[:-1] + ((),))) if len(base_pspec) == 4: - lora_b_pspec = jax.sharding.PartitionSpec((), base_pspec[1], base_pspec[-1]) + lora_b_pspec = jax.sharding.PartitionSpec((), base_pspec[1], + base_pspec[-1]) else: lora_b_pspec = jax.sharding.PartitionSpec((), base_pspec[-1]) else: @@ -874,8 +1026,12 @@ def get_lora_param_sharding(base_param_sharding, lora_module): mesh = base_param_sharding.mesh mem_kind = base_param_sharding.memory_kind return ( - jax.sharding.NamedSharding(mesh=mesh, spec=lora_a_pspec, memory_kind=mem_kind), - jax.sharding.NamedSharding(mesh=mesh, spec=lora_b_pspec, memory_kind=mem_kind), + jax.sharding.NamedSharding(mesh=mesh, + spec=lora_a_pspec, + memory_kind=mem_kind), + jax.sharding.NamedSharding(mesh=mesh, + spec=lora_b_pspec, + memory_kind=mem_kind), ) def module_is_target(module_path): @@ -891,15 +1047,22 @@ def add_lora(out_node, base_node, path): add_lora(out_node[name], child, f"{path}.{name}") else: if name not in ("kernel", "scale", "embedding"): - raise ValueError(f"Unexpected key={name} in base abstract params at {path}") + raise ValueError( + f"Unexpected key={name} in base abstract params at {path}") if not isinstance(child, jax.ShapeDtypeStruct): - raise ValueError(f"Unexpected leaf type {type(child).__name__} at {path}.{name}") + raise ValueError( + f"Unexpected leaf type {type(child).__name__} at {path}.{name}") target_module = module_is_target(path) if target_module is not None: a_shape, b_shape = get_lora_param_shape(child.shape, target_module) - a_sharding, b_sharding = get_lora_param_sharding(child.sharding, target_module) - out_node["lora_a.kernel"] = jax.ShapeDtypeStruct(shape=a_shape, dtype=child.dtype, sharding=a_sharding) - out_node["lora_b.kernel"] = jax.ShapeDtypeStruct(shape=b_shape, dtype=child.dtype, sharding=b_sharding) + a_sharding, b_sharding = get_lora_param_sharding( + child.sharding, target_module) + out_node["lora_a.kernel"] = jax.ShapeDtypeStruct(shape=a_shape, + dtype=child.dtype, + sharding=a_sharding) + out_node["lora_b.kernel"] = jax.ShapeDtypeStruct(shape=b_shape, + dtype=child.dtype, + sharding=b_sharding) else: out_node[name] = None @@ -907,16 +1070,41 @@ def add_lora(out_node, base_node, path): add_lora(lora_abstract_params, base_abstract_params, "") unboxed_abstract_lora_state = train_state.TrainState( - step=0, apply_fn=None, params=lora_abstract_params, tx=None, opt_state={} # type: ignore + step=0, + apply_fn=None, + params=lora_abstract_params, + tx=None, + opt_state={} # type: ignore ) lora_state_mesh_annotations = train_state.TrainState( step=0, apply_fn=None, params=jax.tree_util.tree_map( - lambda x: x.sharding.spec if x.sharding is not None else None, + lambda x: x.sharding.spec + if getattr(x, "sharding", None) is not None else None, lora_abstract_params, ), tx=None, # type: ignore opt_state={}, ) return unboxed_abstract_lora_state, lora_state_mesh_annotations + + +def restore_qlora_base_weights(val): + """Restores qwix custom quantized types from nnx.State representation into unquantized abstract specs.""" + if isinstance(val, (nnx.State, dict)): + pure_dict = {k: restore_qlora_base_weights(v) for k, v in val.items()} + if "array" in pure_dict: + return pure_dict["array"] + if "qvalue" in pure_dict and "scale" in pure_dict: + qval = pure_dict["qvalue"] + scale = pure_dict["scale"] + dtype = getattr(scale, "dtype", jnp.float32) + sharding = getattr(qval, "sharding", getattr(scale, "sharding", None)) + return jax.ShapeDtypeStruct(shape=qval.shape, + dtype=dtype, + sharding=sharding) + return pure_dict + if isinstance(val, nnx.Variable): + return restore_qlora_base_weights(val.get_value()) + return val diff --git a/src/maxtext/utils/max_utils.py b/src/maxtext/utils/max_utils.py index e6aebce059..10e7510835 100644 --- a/src/maxtext/utils/max_utils.py +++ b/src/maxtext/utils/max_utils.py @@ -11,7 +11,6 @@ # 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. - """Common Max Utils needed by multiple modules. All the functions include MaxText modules, such as Pyconfig, should be moved to MaxText utils file. """ @@ -79,7 +78,9 @@ def parse_libtpu_flags_to_dict(flags_str: str) -> dict: match = token_pattern.match(token) if not match: # Throw an error immediately if any token fails the strict format - raise ValueError(f"Invalid flag format detected: '{token}'. Expected format: '--key=value'") + raise ValueError( + f"Invalid flag format detected: '{token}'. Expected format: '--key=value'" + ) key, value = match.groups() @@ -93,15 +94,18 @@ def parse_libtpu_flags_to_dict(flags_str: str) -> dict: def with_memory_kind(t, memory_kind): - return jax.tree_util.tree_map(lambda x: x.with_memory_kind(kind=memory_kind), t) + return jax.tree_util.tree_map(lambda x: x.with_memory_kind(kind=memory_kind), + t) def cast_dtype_from_to(nest, src, dst): """All items in nest with dtype src are casted to dtype dst.""" - return jax.tree_util.tree_map(lambda t: t.astype(dst) if t.dtype == src else t, nest) + return jax.tree_util.tree_map( + lambda t: t.astype(dst) if t.dtype == src else t, nest) def find_nans_and_infs(pytree): + def finder(x): return jnp.any(jnp.isinf(x) | jnp.isnan(x)) @@ -111,7 +115,22 @@ def finder(x): def l2norm_pytree(x): """L2 norm of a pytree of arrays.""" - return jnp.sqrt(jax.tree_util.tree_reduce(lambda x, y: x + jnp.sum(jnp.square(y)), x, initializer=0.0)) + + def _leaf_sq_sum(y): + # Safely skip non-numeric or quantized leaves (e.g., QArray, packed void dtypes, or metadata) + try: + dtype = getattr(y, "dtype", None) + if dtype is not None and (jnp.issubdtype(dtype, jnp.inexact) or + jnp.issubdtype(dtype, jnp.number)): + return jnp.sum(jnp.square(y.astype(jnp.float32))) + except Exception: # pylint: disable=broad-exception-caught + pass + + return 0.0 + + sq_sums = jax.tree_util.tree_map(_leaf_sq_sum, x) + return jnp.sqrt( + jax.tree_util.tree_reduce(lambda a, b: a + b, sq_sums, initializer=0.0)) def calculate_num_params_from_pytree(params): @@ -137,8 +156,10 @@ def calculate_leaf_params_per_chip(arr): shard = arr.addressable_shards[0] return np.prod(shard.data.shape) - params_sizes_per_chip = jax.tree_util.tree_map(calculate_leaf_params_per_chip, params) - total_parameters_per_chip = jax.tree_util.tree_reduce(lambda x, y: x + y, params_sizes_per_chip) + params_sizes_per_chip = jax.tree_util.tree_map(calculate_leaf_params_per_chip, + params) + total_parameters_per_chip = jax.tree_util.tree_reduce(lambda x, y: x + y, + params_sizes_per_chip) return total_parameters_per_chip @@ -162,7 +183,8 @@ def _bytes_of(x): # None or unsupported leaf types: count as zero bytes. if x is not None: - max_logging.log(f"Unsupported leaf type in calculate_bytes_from_pytree: {type(x)}") + max_logging.log( + f"Unsupported leaf type in calculate_bytes_from_pytree: {type(x)}") return 0 @@ -182,7 +204,9 @@ def summarize_size_from_pytree(params): return num_params, num_bytes, num_bytes / num_params -def initialize_summary_writer(tensorboard_dir, run_name, enable_tensorboard=True): +def initialize_summary_writer(tensorboard_dir, + run_name, + enable_tensorboard=True): """Return a tensorboardX SummaryWriter or a no-op stub. In decoupled mode (no Google Cloud), this prefers a repo-local @@ -204,16 +228,22 @@ def initialize_summary_writer(tensorboard_dir, run_name, enable_tensorboard=True try: repo_tb = Path(__file__).resolve().parents[2] / "local_tensorboard" repo_tb.mkdir(parents=True, exist_ok=True) - summary_writer_path = str(repo_tb / run_name) if run_name else str(repo_tb) - max_logging.log(f"Decoupled: using local tensorboard dir {summary_writer_path}") + summary_writer_path = str(repo_tb / + run_name) if run_name else str(repo_tb) + max_logging.log( + f"Decoupled: using local tensorboard dir {summary_writer_path}") return writer.SummaryWriter(summary_writer_path) except Exception as e: # pylint: disable=broad-exception-caught - max_logging.log(f"Decoupled: failed to use local tensorboard dir: {e}; using no-op SummaryWriter.") + max_logging.log( + f"Decoupled: failed to use local tensorboard dir: {e}; using no-op SummaryWriter." + ) return writer.SummaryWriter() # Check if dir or run_name exists! if not tensorboard_dir or not run_name: - max_logging.log("tensorboard_dir or run_name missing; using no-op SummaryWriter to avoid crash.") + max_logging.log( + "tensorboard_dir or run_name missing; using no-op SummaryWriter to avoid crash." + ) return writer.SummaryWriter() summary_writer_path = os.path.join(tensorboard_dir, run_name) @@ -242,18 +272,26 @@ def maybe_initialize_jax_distributed_system(raw_keys): # Early exit for cases where we don't need to initialize the jax distributed system. if raw_keys["skip_jax_distributed_system"]: - max_logging.log("Skipping jax distributed system due to skip_jax_distributed_system=True flag.") + max_logging.log( + "Skipping jax distributed system due to skip_jax_distributed_system=True flag." + ) return if raw_keys["enable_single_controller"]: - max_logging.log("Skipping jax distributed system since its not needed for single controller.") + max_logging.log( + "Skipping jax distributed system since its not needed for single controller." + ) if raw_keys["enable_multi_tier_checkpointing"]: - max_logging.log("Initializing multi-tier checkpointing for single controller...") - mtc_init_kwargs = elastic_utils.single_controller_mtc_init_kwargs(raw_keys) + max_logging.log( + "Initializing multi-tier checkpointing for single controller...") + mtc_init_kwargs = elastic_utils.single_controller_mtc_init_kwargs( + raw_keys) initialize_multi_tier_checkpointing( local_checkpoint_directory=raw_keys["local_checkpoint_directory"], - backup_interval_minutes=raw_keys["multi_tier_checkpointing_backup_interval_minutes"], + backup_interval_minutes=raw_keys[ + "multi_tier_checkpointing_backup_interval_minutes"], run_name=raw_keys["run_name"], - jax_initialization_timeout_seconds=raw_keys["jax_distributed_initialization_timeout"], + jax_initialization_timeout_seconds=raw_keys[ + "jax_distributed_initialization_timeout"], use_colocated_python=True, **mtc_init_kwargs, ) @@ -267,46 +305,63 @@ def maybe_initialize_jax_distributed_system(raw_keys): # Initialization for gpu backend if is_gpu_backend(raw_keys): - max_logging.log("Attempting to initialize the jax distributed system for GPU backend...") + max_logging.log( + "Attempting to initialize the jax distributed system for GPU backend..." + ) initialize_jax_for_gpu(raw_keys) max_logging.log("Jax distributed system initialized on GPU!") return # Initialization for cpu backend if is_cpu_backend(raw_keys): - max_logging.log("Attempting to initialize the jax distributed system for CPU backend...") + max_logging.log( + "Attempting to initialize the jax distributed system for CPU backend..." + ) initialize_jax_for_cpu(raw_keys) max_logging.log("Jax distributed system initialized on CPUs!") return # Initialization for gpu_multiprocess hardware if raw_keys["hardware"] == "gpu_multiprocess": - max_logging.log("Attempting to initialize the jax distributed system for gpu_multiprocess hardware...") + max_logging.log( + "Attempting to initialize the jax distributed system for gpu_multiprocess hardware..." + ) if not raw_keys["enable_emergency_checkpoint"]: - jax.distributed.initialize(initialization_timeout=raw_keys["jax_distributed_initialization_timeout"]) + jax.distributed.initialize(initialization_timeout=raw_keys[ + "jax_distributed_initialization_timeout"]) else: - max_logging.log("Initializing jax distributed to support local checkpointing with GPUs...") - jax.distributed.initialize(initialization_timeout=raw_keys["jax_distributed_initialization_timeout"]) + max_logging.log( + "Initializing jax distributed to support local checkpointing with GPUs..." + ) + jax.distributed.initialize(initialization_timeout=raw_keys[ + "jax_distributed_initialization_timeout"]) ocp.multihost.initialize_runtime_to_distributed_ids() ocp.multihost.initialize_distributed_to_device_ids() max_logging.log("Jax distributed system initialized!") return # Initialization for tpu backend - max_logging.log("Attempting to initialize the jax distributed system for TPU backend...") + max_logging.log( + "Attempting to initialize the jax distributed system for TPU backend...") if raw_keys["enable_multi_tier_checkpointing"]: initialize_multi_tier_checkpointing( local_checkpoint_directory=raw_keys["local_checkpoint_directory"], - backup_interval_minutes=raw_keys["multi_tier_checkpointing_backup_interval_minutes"], + backup_interval_minutes=raw_keys[ + "multi_tier_checkpointing_backup_interval_minutes"], run_name=raw_keys["run_name"], - jax_initialization_timeout_seconds=raw_keys["jax_distributed_initialization_timeout"], + jax_initialization_timeout_seconds=raw_keys[ + "jax_distributed_initialization_timeout"], data_parallelism=raw_keys["mtc_data_parallelism"], num_slices=raw_keys["num_slices"], ) - max_logging.log("Jax distributed system initialized on TPUs for multi-tier checkpointing!") - elif raw_keys["enable_checkpointing"] and raw_keys["compile_topology_num_slices"] == -1: + max_logging.log( + "Jax distributed system initialized on TPUs for multi-tier checkpointing!" + ) + elif raw_keys["enable_checkpointing"] and raw_keys[ + "compile_topology_num_slices"] == -1: if not raw_keys["enable_emergency_checkpoint"]: - jax.distributed.initialize(initialization_timeout=raw_keys["jax_distributed_initialization_timeout"]) + jax.distributed.initialize(initialization_timeout=raw_keys[ + "jax_distributed_initialization_timeout"]) else: initialize_jax_for_tpu_with_emergency_checkpointing(raw_keys) max_logging.log("Jax distributed system initialized on TPUs!") @@ -321,7 +376,8 @@ def initialize_jax_for_gpu(raw_keys): for env_var in env_var_list: devices = os.getenv(env_var) if devices is not None: - max_logging.log(f"Using {env_var} to initialize JAX distributed system: {devices}") + max_logging.log( + f"Using {env_var} to initialize JAX distributed system: {devices}") break if devices is None: jax.config.update("jax_cuda_visible_devices", "all") @@ -336,9 +392,12 @@ def initialize_jax_for_gpu(raw_keys): jax.distributed.initialize( coordinator_address=f"{coordinator_ip}:{coordinator_port}", - num_processes=int(os.getenv("NNODES")), # pyrefly: ignore[bad-argument-type] - process_id=int(os.getenv("NODE_RANK")), # pyrefly: ignore[bad-argument-type] - initialization_timeout=raw_keys["jax_distributed_initialization_timeout"], + num_processes=int( + os.getenv("NNODES")), # pyrefly: ignore[bad-argument-type] + process_id=int( + os.getenv("NODE_RANK")), # pyrefly: ignore[bad-argument-type] + initialization_timeout=raw_keys[ + "jax_distributed_initialization_timeout"], local_device_ids=devices, ) max_logging.log(f"JAX global devices: {jax.devices()}") @@ -349,16 +408,20 @@ def initialize_jax_for_cpu(raw_keys): coordinator_ip_address = get_coordinator_ip_address() coordinator_address = coordinator_ip_address + ":1234" # JAX coordinator port used in XPK # Env variables to be set in XPK or otherwise - job_index = int(os.environ.get("JOB_INDEX")) # pyrefly: ignore[bad-argument-type] - job_completion_index = int(os.environ.get("JOB_COMPLETION_INDEX")) # pyrefly: ignore[bad-argument-type] - processes_in_job = int(os.environ.get("PROCESSES_IN_JOB")) # pyrefly: ignore[bad-argument-type] + job_index = int( + os.environ.get("JOB_INDEX")) # pyrefly: ignore[bad-argument-type] + job_completion_index = int(os.environ.get( + "JOB_COMPLETION_INDEX")) # pyrefly: ignore[bad-argument-type] + processes_in_job = int( + os.environ.get("PROCESSES_IN_JOB")) # pyrefly: ignore[bad-argument-type] pid = job_index * processes_in_job + job_completion_index max_logging.log(f" Jax process id is {pid} ") # Explicit initialize is needed only for CPUs jax.distributed.initialize( coordinator_address=coordinator_address, process_id=pid, - num_processes=int(os.environ.get("JAX_PROCESS_COUNT")), # pyrefly: ignore[bad-argument-type] + num_processes=int(os.environ.get( + "JAX_PROCESS_COUNT")), # pyrefly: ignore[bad-argument-type] initialization_timeout=raw_keys["jax_distributed_initialization_timeout"], ) @@ -374,12 +437,12 @@ def initialize_jax_for_tpu_with_emergency_checkpointing(raw_keys): if process_id != "" and coordinator_address != "": max_logging.log( f"Using {process_id} as the process_id and {coordinator_address} as the" - " coordinator_address to initialize JAX distributed runtime..." - ) + " coordinator_address to initialize JAX distributed runtime...") jax.distributed.initialize( coordinator_address=coordinator_address, process_id=int(process_id), - initialization_timeout=raw_keys["jax_distributed_initialization_timeout"], + initialization_timeout=raw_keys[ + "jax_distributed_initialization_timeout"], ) ocp.multihost.initialize_runtime_to_distributed_ids() @@ -389,7 +452,8 @@ def initialize_jax_for_tpu_with_emergency_checkpointing(raw_keys): def _retrieve_jax_init_info(raw_keys): """Retrieve JAX init info from a local file.""" JAX_INIT_INFO_FILE = "jax-init-info.txt" - local_jax_init_info_file = epath.Path(raw_keys["local_checkpoint_directory"]) / JAX_INIT_INFO_FILE + local_jax_init_info_file = epath.Path( + raw_keys["local_checkpoint_directory"]) / JAX_INIT_INFO_FILE # Allow time for the JAX init info file to be populated by GKE. This is needed because the file is # only populated when the worker with process id of 0 is determined. After a disruption, although some # workers might be up and running, the init info file won't be populated until the node with process id @@ -398,18 +462,20 @@ def _retrieve_jax_init_info(raw_keys): for i in range(900): if local_jax_init_info_file.exists(): return local_jax_init_info_file.read_text().split("\n")[:2] - max_logging.log(f"Unable to locate {JAX_INIT_INFO_FILE} after {i} seconds, sleeping for 1 second before retrying...") + max_logging.log( + f"Unable to locate {JAX_INIT_INFO_FILE} after {i} seconds, sleeping for 1 second before retrying..." + ) time.sleep(1) - max_logging.log( - f"Unable to locate {JAX_INIT_INFO_FILE} after 900 seconds," "returning empty process id and coordinator address." - ) + max_logging.log(f"Unable to locate {JAX_INIT_INFO_FILE} after 900 seconds," + "returning empty process id and coordinator address.") return "", "" def get_num_slices(raw_keys, config=None): """Calculate num_slices based on number of devices.""" if raw_keys.get("num_slices", -1) != -1: - max_logging.log(f"Using num_slices={raw_keys['num_slices']} per user request.") + max_logging.log( + f"Using num_slices={raw_keys['num_slices']} per user request.") return raw_keys["num_slices"] if getattr(raw_keys, "hardware", None) == "cpu": max_logging.log(" Setting num_slices=1 for CPU hardware type") @@ -444,7 +510,8 @@ def get_coordinator_ip_address(): max_coordinator_lookups = 50 while not coordinator_found and lookup_attempt <= max_coordinator_lookups: try: - coordinator_ip_address = socket.gethostbyname(coordinator_address) # pyrefly: ignore[bad-argument-type] + coordinator_ip_address = socket.gethostbyname( + coordinator_address) # pyrefly: ignore[bad-argument-type] coordinator_found = True except socket.gaierror: max_logging.log( @@ -456,7 +523,8 @@ def get_coordinator_ip_address(): return coordinator_ip_address -def fill_unspecified_mesh_axes(parallelism_vals, target_product, parallelism_type): +def fill_unspecified_mesh_axes(parallelism_vals, target_product, + parallelism_type): """Evaluates unspecified DCN/ICI parallelism values""" if -1 in parallelism_vals: assert ( @@ -466,9 +534,8 @@ def fill_unspecified_mesh_axes(parallelism_vals, target_product, parallelism_typ determined_val = target_product / np.prod(parallelism_vals) * -1 - assert ( - determined_val >= 1 and determined_val.is_integer() - ), f"Unspecified value unable to be determined with the given\ + assert (determined_val >= 1 and determined_val.is_integer() + ), f"Unspecified value unable to be determined with the given\ {parallelism_type} parallelism values" parallelism_vals[parallelism_vals.index(-1)] = int(determined_val) @@ -492,7 +559,9 @@ def reshape_mesh_to_rings(a, strategy): a_i = i * 2 a_j = j * 2 # forms a ring of size 4 - b[i].append([a[a_i, a_j], a[a_i, a_j + 1], a[a_i + 1, a_j + 1], a[a_i + 1, a_j]]) + b[i].append([ + a[a_i, a_j], a[a_i, a_j + 1], a[a_i + 1, a_j + 1], a[a_i + 1, a_j] + ]) b = np.array(b) b = np.reshape(b, (64, 4)) elif strategy == HYBRID_RING_32X8: @@ -502,22 +571,21 @@ def reshape_mesh_to_rings(a, strategy): a_i = i * 2 a_j = j * 4 # forms a ring of size 8 - b[i].append( - [ - a[a_i, a_j], - a[a_i, a_j + 1], - a[a_i, a_j + 2], - a[a_i, a_j + 3], - a[a_i + 1, a_j + 3], - a[a_i + 1, a_j + 2], - a[a_i + 1, a_j + 1], - a[a_i + 1, a_j], - ] - ) + b[i].append([ + a[a_i, a_j], + a[a_i, a_j + 1], + a[a_i, a_j + 2], + a[a_i, a_j + 3], + a[a_i + 1, a_j + 3], + a[a_i + 1, a_j + 2], + a[a_i + 1, a_j + 1], + a[a_i + 1, a_j], + ]) b = np.array(b) b = np.reshape(b, (32, 8)) else: - raise ValueError(f"The strategy {strategy} to reshape the mesh is not implemented.") + raise ValueError( + f"The strategy {strategy} to reshape the mesh is not implemented.") return b @@ -530,31 +598,38 @@ def create_custom_device_mesh( should_sort_granules_by_key: bool = True, ) -> np.ndarray: """Custom device mesh for 64x4 ici parallelism""" - assert len(devices) % 256 == 0, f"This custom mesh is not valid for {len(devices)} devices" + assert len( + devices + ) % 256 == 0, f"This custom mesh is not valid for {len(devices)} devices" attr = "process_index" if process_is_granule else "slice_index" if not hasattr(devices[0], attr): - raise ValueError(f"Device {devices[0]} does not have attribute {attr}. See" " `process_is_granule` option.") + raise ValueError(f"Device {devices[0]} does not have attribute {attr}. See" + " `process_is_granule` option.") granule_dict = collections.defaultdict(list) for dev in devices: granule_dict[getattr(dev, attr)].append(dev) - granules = ( - [granule_dict[key] for key in sorted(granule_dict.keys())] if should_sort_granules_by_key else granule_dict.values() - ) + granules = ([granule_dict[key] for key in sorted(granule_dict.keys())] + if should_sort_granules_by_key else granule_dict.values()) if np.prod(dcn_mesh_shape) != len(granules): - raise ValueError(f"Number of slices {len(granules)} must equal the product of " f"dcn_mesh_shape {dcn_mesh_shape}") + raise ValueError( + f"Number of slices {len(granules)} must equal the product of " + f"dcn_mesh_shape {dcn_mesh_shape}") per_granule_meshes = [ mesh_utils.create_device_mesh( [16, 16], granule, allow_split_physical_axes=False, - ) - for granule in granules + ) for granule in granules ] - per_granule_meshes = [np.reshape(reshape_mesh_to_rings(x, custom_strategy), mesh_shape) for x in per_granule_meshes] + per_granule_meshes = [ + np.reshape(reshape_mesh_to_rings(x, custom_strategy), mesh_shape) + for x in per_granule_meshes + ] # TODO(jekbradbury): handle non-uniform DCN topologies granule_mesh = np.arange(len(granules)).reshape(dcn_mesh_shape) - blocks = np.vectorize(lambda i: per_granule_meshes[i], otypes=[object])(granule_mesh) + blocks = np.vectorize(lambda i: per_granule_meshes[i], + otypes=[object])(granule_mesh) device_mesh = np.block(blocks.tolist()) return device_mesh @@ -573,7 +648,9 @@ def is_valid_custom_mesh(ici_parallelism, strategy): if sorted(set(ici_parallelism)) == valid_strategies[strategy]: return True else: - raise ValueError(f"Invalid custom_mesh:{strategy} chosen for ICI mesh shape {ici_parallelism}") + raise ValueError( + f"Invalid custom_mesh:{strategy} chosen for ICI mesh shape {ici_parallelism}" + ) else: raise ValueError(f"The strategy {strategy} to reshape the mesh is invalid.") @@ -589,8 +666,10 @@ def optimize_mesh_for_tpu_v6e(mesh, devices): # check that the physical topology is 2x4 device_coords = [d.coords for d in devices] coord_size = len(device_coords[0]) - max_coords = tuple(max(dc[i] for dc in device_coords) for i in range(coord_size)) - min_coords = tuple(min(dc[i] for dc in device_coords) for i in range(coord_size)) + max_coords = tuple( + max(dc[i] for dc in device_coords) for i in range(coord_size)) + min_coords = tuple( + min(dc[i] for dc in device_coords) for i in range(coord_size)) dims = tuple(h - l + 1 for (h, l) in zip(max_coords, min_coords)) if dims != (2, 4, 1): return mesh @@ -612,7 +691,8 @@ def unbox_logicallypartioned(boxed_pytree): a pytree where all all LogicallyPartitioned leaves have been unboxed. """ return jax.tree_util.tree_map( - lambda x: x.unbox() if isinstance(x, flax.linen.spmd.LogicallyPartitioned) else x, + lambda x: x.unbox() + if isinstance(x, flax.linen.spmd.LogicallyPartitioned) else x, boxed_pytree, is_leaf=lambda k: isinstance(k, flax.linen.spmd.LogicallyPartitioned), ) @@ -654,7 +734,11 @@ def cross_entropy_with_logits( return loss, total_z_loss -def _cross_entropy_with_logits_fwd(logits: jnp.ndarray, targets: jnp.ndarray, z_loss: float = 0.0) -> tuple[ +def _cross_entropy_with_logits_fwd( + logits: jnp.ndarray, + targets: jnp.ndarray, + z_loss: float = 0.0 +) -> tuple[ tuple[jnp.ndarray, jnp.ndarray], tuple[ jnp.ndarray, @@ -698,11 +782,14 @@ def _cross_entropy_with_logits_bwd( g: tuple[jnp.ndarray, jnp.ndarray], ) -> tuple[jnp.ndarray, None, None]: """Backward-mode of `cross_entropy_with_logits`.""" - g = g[0] # Ignore z_loss component as that is only used for logging. # pyrefly: ignore[bad-assignment] + g = g[ + 0] # Ignore z_loss component as that is only used for logging. # pyrefly: ignore[bad-assignment] logits, targets, z_loss, exp_shifted, sum_exp, log_z = res # z-loss term adds the (2 * z_loss * log_z) factor. - deriv = jnp.expand_dims(1 + 2 * z_loss * log_z, -1) * exp_shifted / sum_exp - targets - g_logits = jnp.expand_dims(g, axis=-1) * deriv # pyrefly: ignore[bad-argument-type] + deriv = jnp.expand_dims(1 + 2 * z_loss * log_z, + -1) * exp_shifted / sum_exp - targets + g_logits = jnp.expand_dims( + g, axis=-1) * deriv # pyrefly: ignore[bad-argument-type] return ( jnp.asarray(g_logits, logits.dtype), @@ -711,7 +798,8 @@ def _cross_entropy_with_logits_bwd( ) # sets z-loss coeff gradient to 0 -cross_entropy_with_logits.defvjp(_cross_entropy_with_logits_fwd, _cross_entropy_with_logits_bwd) +cross_entropy_with_logits.defvjp(_cross_entropy_with_logits_fwd, + _cross_entropy_with_logits_bwd) def print_pytree_shape(print_str, ptree): @@ -731,18 +819,24 @@ def get_project(): if is_decoupled(): return os.environ.get("LOCAL_GCLOUD_PROJECT", "local-maxtext-project") try: - completed_command = subprocess.run(["gcloud", "config", "get", "project"], check=True, capture_output=True) + completed_command = subprocess.run(["gcloud", "config", "get", "project"], + check=True, + capture_output=True) project_outputs = completed_command.stdout.decode().strip().split("\n") if len(project_outputs) < 1 or project_outputs[-1] == "": - max_logging.log("You must specify config.vertex_tensorboard_project or set 'gcloud config set project '") + max_logging.log( + "You must specify config.vertex_tensorboard_project or set 'gcloud config set project '" + ) return None return project_outputs[-1] except (FileNotFoundError, subprocess.CalledProcessError) as ex: - max_logging.log(f"Unable to retrieve gcloud project (decoupled={is_decoupled()}): {ex}") + max_logging.log( + f"Unable to retrieve gcloud project (decoupled={is_decoupled()}): {ex}") return None def delete_pytree(p): + def delete_leaf(leaf): if isinstance(leaf, jax.Array): leaf.delete() @@ -753,23 +847,20 @@ def delete_leaf(leaf): def summarize_pytree_data(params, name="Params", raw=False): """Generate basic metrics of a given Pytree.""" - num_params, total_param_size, avg_param_size = summarize_size_from_pytree(params) + num_params, total_param_size, avg_param_size = summarize_size_from_pytree( + params) if not raw: num_params_in_billions = num_params / 1e9 total_param_size_in_gb = total_param_size / 1e9 - print( - f"{name} stats: \n" - f"\tTotal number of params: {num_params_in_billions:.3f} billion \n" - f"\tTotal memory usage: {total_param_size_in_gb:.3f} GB \n" - f"\tAvg size: {avg_param_size:.3f} bytes\n" - ) + print(f"{name} stats: \n" + f"\tTotal number of params: {num_params_in_billions:.3f} billion \n" + f"\tTotal memory usage: {total_param_size_in_gb:.3f} GB \n" + f"\tAvg size: {avg_param_size:.3f} bytes\n") else: - print( - f"{name} stats: \n" - f"\tTotal number of params: {num_params:.3f} \n" - f"\tTotal memory usage: {total_param_size:.3f} bytes \n" - f"\tAvg size: {avg_param_size:.3f} bytes\n" - ) + print(f"{name} stats: \n" + f"\tTotal number of params: {num_params:.3f} \n" + f"\tTotal memory usage: {total_param_size:.3f} bytes \n" + f"\tAvg size: {avg_param_size:.3f} bytes\n") return num_params, total_param_size, avg_param_size @@ -795,7 +886,9 @@ def print_cpu_ram_stats(label: str): available = round(ram.available / 2**30, 2) used = round(ram.used / 2**30, 2) - max_logging.log(f"\tUsing (GB) {used} / {total} ({used/total:%}) --> Available:{available}") + max_logging.log( + f"\tUsing (GB) {used} / {total} ({used/total:%}) --> Available:{available}" + ) except (RuntimeError, KeyError, TypeError) as ex: max_logging.log(f"\tRAM stats unavailable, error: {ex}") @@ -819,8 +912,7 @@ def bytes_to_gb(num_bytes): f"Total estimated memory size: {total_gb:.1f} GB, estimated output" f" size: {output_gb:.1f} GB, estimated temp size: {temp_gb:.1f} GB, " f"estimated argument size: {argument_gb:.1f} GB, Estimated host temp" - f" size: {host_temp_gb:.1f} GB." - ) + f" size: {host_temp_gb:.1f} GB.") max_logging.log("Note that compiler could over-estimate the HBM usage.") @@ -829,7 +921,9 @@ def print_system_information(): Note that this will initialize the JAX backend.""" max_logging.log(f"System Information: Jax Version: {jax.__version__}") max_logging.log(f"System Information: Jaxlib Version: {jax.lib.__version__}") - max_logging.log(f"System Information: Jax Backend: {jax.extend.backend.get_backend().platform_version}") + max_logging.log( + f"System Information: Jax Backend: {jax.extend.backend.get_backend().platform_version}" + ) def permute_to_match_maxtext_rope(arr): @@ -854,7 +948,10 @@ def unpermute_from_match_maxtext_rope(arr, model_size): @partial(jax.jit, static_argnames=("cp_size", "seq_dim", "to_contiguous")) -def reorder_sequence(tensor, cp_size: int, seq_dim: int = 1, to_contiguous: bool = False): +def reorder_sequence(tensor, + cp_size: int, + seq_dim: int = 1, + to_contiguous: bool = False): """Reorders the sequence of the tensor. For example, with cp_size=2, [0, 1, 2, 3, 4, 5, 6, 7] -> [0, 1, 6, 7, 2, 3, 4, 5] and backward @@ -885,7 +982,7 @@ def reorder_sequence(tensor, cp_size: int, seq_dim: int = 1, to_contiguous: bool *ori_tensor_shape[:seq_dim], 2 * cp_size, group_size, - *ori_tensor_shape[seq_dim + 1 :], + *ori_tensor_shape[seq_dim + 1:], ) # Swap target seq_dim with axis 0 to perform slicing/concat easily: @@ -915,7 +1012,10 @@ def reorder_sequence(tensor, cp_size: int, seq_dim: int = 1, to_contiguous: bool @partial(jax.jit, static_argnums=(1, 2, 3)) -def reorder_causal_load_balanced(batch, cp_size, reorder_strategy, hardware="tpu"): +def reorder_causal_load_balanced(batch, + cp_size, + reorder_strategy, + hardware="tpu"): """Reorders the example batch sequences using a hardware-appropriate backend. On GPU (hardware="gpu" or "gpu_multiprocess"), uses Transformer Engine's @@ -969,14 +1069,13 @@ def reorder_causal_load_balanced(batch, cp_size, reorder_strategy, hardware="tpu } return { - key: reorder_causal_load_balancing( - value, - reorder_strategy_map[reorder_strategy], - cp_size=cp_size, - seq_dim=1, - ) - if key in _reorder_keys - else value + key: + reorder_causal_load_balancing( + value, + reorder_strategy_map[reorder_strategy], + cp_size=cp_size, + seq_dim=1, + ) if key in _reorder_keys else value for key, value in batch.items() } else: @@ -985,12 +1084,11 @@ def reorder_causal_load_balanced(batch, cp_size, reorder_strategy, hardware="tpu f"STRIPED reorder strategy requires Transformer Engine and is only supported on GPU, got hardware={hardware!r}." ) return { - key: reorder_sequence( - value, - cp_size=cp_size, - ) - if key in _reorder_keys - else value + key: + reorder_sequence( + value, + cp_size=cp_size, + ) if key in _reorder_keys else value for key, value in batch.items() } @@ -1027,7 +1125,7 @@ def reorder_mask_load_balancing(tensor, cp_size: int, seq_dim: int): *ori_tensor_shape[:seq_dim], 2 * cp_size, group_size, - *ori_tensor_shape[seq_dim + 1 :], + *ori_tensor_shape[seq_dim + 1:], ) # Create first and second halves @@ -1080,7 +1178,7 @@ def unscan_train_state_params(params, sharding, mesh, scan_axis, layer_groups): def strip_scan_axis(pspec: P) -> P: """Removes the element at `scan_axis` from a PartitionSpec tuple.""" spec_tuple = tuple(pspec) - return P(*(spec_tuple[:scan_axis] + spec_tuple[scan_axis + 1 :])) + return P(*(spec_tuple[:scan_axis] + spec_tuple[scan_axis + 1:])) for layer_name, num_layers in layer_groups: scanned_layers = decoder[layer_name] @@ -1089,15 +1187,18 @@ def strip_scan_axis(pspec: P) -> P: # 1. Compute the target sharding for a single, unscanned layer. # This is done once per layer group, with no expensive compilation. unscanned_sharding_spec = jax.tree_util.tree_map( - strip_scan_axis, jax.tree_util.tree_map(lambda x: x.spec, scanned_sharding) - ) - unscanned_sharding = jax.tree_util.tree_map(lambda ps: jax.sharding.NamedSharding(mesh, ps), unscanned_sharding_spec) + strip_scan_axis, + jax.tree_util.tree_map(lambda x: x.spec, scanned_sharding)) + unscanned_sharding = jax.tree_util.tree_map( + lambda ps: jax.sharding.NamedSharding(mesh, ps), + unscanned_sharding_spec) # 2. Create a list of PyTrees, one for each layer, by slicing the original. # This is more direct than repeatedly calling a JIT'd function. layer_pytrees = [ - jax.tree_util.tree_map(functools.partial(jnp.take, indices=i, axis=scan_axis), scanned_layers) - for i in range(num_layers) + jax.tree_util.tree_map( + functools.partial(jnp.take, indices=i, axis=scan_axis), + scanned_layers) for i in range(num_layers) ] # 3. Reshard each layer's PyTree and assign it to the new key. @@ -1109,7 +1210,8 @@ def strip_scan_axis(pspec: P) -> P: del decoder[layer_name] -def rescan_train_state_params(params, source_shardings, scan_axis, layer_groups): +def rescan_train_state_params(params, source_shardings, scan_axis, + layer_groups): """ Reconstruct scanned layers from per-layer entries using minimal HBM. @@ -1125,7 +1227,8 @@ def rescan_train_state_params(params, source_shardings, scan_axis, layer_groups) for layer_name, num_layers in layer_groups: def stack_layers(*layers): - return jax.tree_util.tree_map(lambda *xs: jnp.stack(xs, axis=scan_axis), *layers) + return jax.tree_util.tree_map(lambda *xs: jnp.stack(xs, axis=scan_axis), + *layers) # Create a wrapper that allows pjit + donation compiled_stack = jax.jit( diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index c7530910eb..2198e5c893 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -63,8 +63,7 @@ def get_input_data_sharding(config, mesh): def assert_params_sufficiently_sharded(params, mesh, tolerance): max_logging.log( "WARNING: Function maxtext_utils.assert_params_sufficiently_sharded is deprecated." - "Please use sharding.assert_params_sufficiently_sharded." - ) + "Please use sharding.assert_params_sufficiently_sharded.") return sharding.assert_params_sufficiently_sharded(params, mesh, tolerance) @@ -78,61 +77,80 @@ def add_data_to_sharding(mesh, path, aval, shardings): def maybe_update_params_sharding_with_opt(config, state_mesh_shardings): max_logging.log( "WARNING: Function maxtext_utils.maybe_update_params_sharding_with_opt is deprecated." - "Please use sharding.maybe_update_params_sharding_with_opt." - ) - return sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings) + "Please use sharding.maybe_update_params_sharding_with_opt.") + return sharding.maybe_update_params_sharding_with_opt(config, + state_mesh_shardings) -def all_gather_over_fsdp(variables, sharding_info, mesh, logical_axis_rules, shard_mode): +def all_gather_over_fsdp(variables, sharding_info, mesh, logical_axis_rules, + shard_mode): max_logging.log( "WARNING: Function maxtext_utils.all_gather_over_fsdp is deprecated. Please use sharding.all_gather_over_fsdp." ) - return sharding.all_gather_over_fsdp(variables, sharding_info, mesh, logical_axis_rules, shard_mode) + return sharding.all_gather_over_fsdp(variables, sharding_info, mesh, + logical_axis_rules, shard_mode) -def get_functional_train_with_signature( - train_step, data_sharding, state_mesh_shardings, model, config, params_shardings=None -): +def get_functional_train_with_signature(train_step, + data_sharding, + state_mesh_shardings, + model, + config, + params_shardings=None): """Get the shardings (both state and data) for `train_step`.""" - functional_train = functools.partial(train_step, model, config, state_mesh_shardings, params_shardings) + functional_train = functools.partial(train_step, model, config, + state_mesh_shardings, params_shardings) functional_train.__name__ = "train_step" # pyrefly: ignore[missing-attribute] if config.pure_nnx: in_shardings = (state_mesh_shardings, data_sharding) # State, batch else: - in_shardings = (state_mesh_shardings, data_sharding, None) # State, batch, rng + in_shardings = (state_mesh_shardings, data_sharding, None + ) # State, batch, rng out_shardings = (state_mesh_shardings, None) # State, metrics static_argnums = () # We partial out the static argnums of model and config donate_argnums = 0 # This is the index of the state - we allow the compiler to make use of this memory. return functional_train, in_shardings, out_shardings, static_argnums, donate_argnums -def get_functional_eval_with_signature(eval_step, data_sharding, state_mesh_shardings, model, config): +def get_functional_eval_with_signature(eval_step, data_sharding, + state_mesh_shardings, model, config): """Get the shardings (both state and data) for `eval_step`.""" functional_eval = functools.partial(eval_step, model, config) functional_eval.__name__ = "eval_step" # pyrefly: ignore[missing-attribute] if config.pure_nnx: - in_shardings = (state_mesh_shardings, data_sharding) # State, batch (NNX: no rng) + in_shardings = (state_mesh_shardings, data_sharding + ) # State, batch (NNX: no rng) else: - in_shardings = (state_mesh_shardings, data_sharding, None) # State, batch, rng + in_shardings = (state_mesh_shardings, data_sharding, None + ) # State, batch, rng out_shardings = None # metrics static_argnums = () # We partial out the static argnums of model, config - donate_argnums = () # state will be kept instead of being donated in eval_step + donate_argnums = ( + ) # state will be kept instead of being donated in eval_step return functional_eval, in_shardings, out_shardings, static_argnums, donate_argnums def shard_reorder_causal_load_balanced( - batch, cp_size, shard_mode, reorder_strategy=ReorderStrategy.DUAL_CHUNK_SWAP, hardware="tpu" -): + batch, + cp_size, + shard_mode, + reorder_strategy=ReorderStrategy.DUAL_CHUNK_SWAP, + hardware="tpu"): """Shard the output of the reordered sequence.""" - reordered = max_utils.reorder_causal_load_balanced(batch, cp_size, reorder_strategy, hardware) + reordered = max_utils.reorder_causal_load_balanced(batch, cp_size, + reorder_strategy, hardware) for _, v in batch.items(): if isinstance(v, jax.Array): - reordered = sharding.maybe_shard_with_name(reordered, v.sharding, shard_mode) + reordered = sharding.maybe_shard_with_name(reordered, v.sharding, + shard_mode) break return reordered -def get_reorder_callable(cp_size, shard_mode, reorder_strategy=ReorderStrategy.DUAL_CHUNK_SWAP, hardware="tpu"): +def get_reorder_callable(cp_size, + shard_mode, + reorder_strategy=ReorderStrategy.DUAL_CHUNK_SWAP, + hardware="tpu"): """Creates a callable that can be used with map() to reorder batches.""" return functools.partial( shard_reorder_causal_load_balanced, @@ -155,21 +173,34 @@ def get_shaped_batch(config, batch_sharding=None): else: batch_shape = (config.global_batch_size_to_load, config.max_target_length) shaped_batch = {} - shaped_batch["inputs"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) - shaped_batch["inputs_position"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) - shaped_batch["inputs_segmentation"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) - shaped_batch["targets"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) - shaped_batch["targets_position"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) - shaped_batch["targets_segmentation"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) + shaped_batch["inputs"] = jax.ShapeDtypeStruct(batch_shape, + jnp.int32, + sharding=batch_sharding) + shaped_batch["inputs_position"] = jax.ShapeDtypeStruct( + batch_shape, jnp.int32, sharding=batch_sharding) + shaped_batch["inputs_segmentation"] = jax.ShapeDtypeStruct( + batch_shape, jnp.int32, sharding=batch_sharding) + shaped_batch["targets"] = jax.ShapeDtypeStruct(batch_shape, + jnp.int32, + sharding=batch_sharding) + shaped_batch["targets_position"] = jax.ShapeDtypeStruct( + batch_shape, jnp.int32, sharding=batch_sharding) + shaped_batch["targets_segmentation"] = jax.ShapeDtypeStruct( + batch_shape, jnp.int32, sharding=batch_sharding) if config.use_multimodal: image_shape = mm_processor.get_dummy_image_shape_for_init( - config.model_name, batch_size=config.micro_batch_size_to_train_on - ) - shaped_batch["images"] = jax.ShapeDtypeStruct(image_shape, jnp.int32, sharding=batch_sharding) - shaped_batch["image_masks"] = jax.ShapeDtypeStruct(image_shape[:2], jnp.int32, sharding=batch_sharding) + config.model_name, batch_size=config.micro_batch_size_to_train_on) + shaped_batch["images"] = jax.ShapeDtypeStruct(image_shape, + jnp.int32, + sharding=batch_sharding) + shaped_batch["image_masks"] = jax.ShapeDtypeStruct(image_shape[:2], + jnp.int32, + sharding=batch_sharding) if config.use_audio: audio_shape = mm_processor.get_dummy_audio_shape_for_init(config) - shaped_batch["audios"] = jax.ShapeDtypeStruct(audio_shape, jnp.float32, sharding=batch_sharding) + shaped_batch["audios"] = jax.ShapeDtypeStruct(audio_shape, + jnp.float32, + sharding=batch_sharding) return shaped_batch @@ -189,7 +220,8 @@ def should_prevent_cse_in_remat(config): if config.scan_layers: return False - if config.gradient_accumulation_steps > 1 and config.hardware in ("gpu", "gpu_multiprocess"): + if config.gradient_accumulation_steps > 1 and config.hardware in ( + "gpu", "gpu_multiprocess"): return False return True @@ -232,7 +264,8 @@ def get_save_and_offload_names(config) -> tuple[list[str], list[str]]: "mlpwo", ] if config.remat_policy == "custom": - return list(config.tensors_on_device or []), list(config.tensors_to_offload or []) + return list(config.tensors_on_device or + []), list(config.tensors_to_offload or []) return [], [] @@ -247,7 +280,8 @@ def load_serialized_compiled(save_name): return f.read() def get_train_input_output_trees(func, input_args, input_kwargs): - _, in_tree_recreated = jax.tree_util.tree_flatten((input_args, input_kwargs)) + _, in_tree_recreated = jax.tree_util.tree_flatten( + (input_args, input_kwargs)) out_shaped = jax.eval_shape(func, *input_args, **input_kwargs) _, out_tree_recreated = jax.tree_util.tree_flatten(out_shaped) return in_tree_recreated, out_tree_recreated @@ -260,8 +294,13 @@ def get_train_input_output_trees(func, input_args, input_kwargs): example_rng = jax.random.PRNGKey(0) shaped_input_args = (state, shaped_batch, example_rng) shaped_input_kwargs = {} - in_tree, out_tree = get_train_input_output_trees(partial_train, shaped_input_args, shaped_input_kwargs) - p_train_step = deserialize_and_load(serialized_compiled, in_tree, out_tree, execution_devices=execution_devices) + in_tree, out_tree = get_train_input_output_trees(partial_train, + shaped_input_args, + shaped_input_kwargs) + p_train_step = deserialize_and_load(serialized_compiled, + in_tree, + out_tree, + execution_devices=execution_devices) return p_train_step @@ -270,36 +309,34 @@ def calculate_tokens_training_per_device(config): return config.max_target_length * config.per_device_batch_size * config.gradient_accumulation_steps -def calculate_gemma2_tflops_training_per_device(config, total_ffn_flops, qkv_flops, projection_flops, embedding_flops): +def calculate_gemma2_tflops_training_per_device(config, total_ffn_flops, + qkv_flops, projection_flops, + embedding_flops): """ Calculate training TFLOP for Gemma2 as in Gemma2 we combine [local_attention, global_attention] into one decoder layer and we use sliding window attention in local_attention """ window = min(config.sliding_window_size, config.max_target_length) - global_causal_flops = ( - 2 * config.per_device_batch_size * config.max_target_length**2 * config.num_query_heads * config.head_dim - ) - local_causal_flops = ( - 4 - * config.per_device_batch_size - * (config.max_target_length * window - 0.5 * window**2) - * config.num_query_heads - * config.head_dim - ) + global_causal_flops = (2 * config.per_device_batch_size * + config.max_target_length**2 * config.num_query_heads * + config.head_dim) + local_causal_flops = (4 * config.per_device_batch_size * + (config.max_target_length * window - 0.5 * window**2) * + config.num_query_heads * config.head_dim) causal_attention_flops = global_causal_flops + local_causal_flops attention_tflops = causal_attention_flops * config.num_decoder_layers * 3 / 10**12 # multiply num_decoder_layers by 2 because we combine [local_attention, global_attention] into one decoder layer - learnable_weight_tflops = ( - ((total_ffn_flops + qkv_flops + projection_flops) * config.num_decoder_layers * 2 + embedding_flops) * 3 / 10**12 - ) + learnable_weight_tflops = (((total_ffn_flops + qkv_flops + projection_flops) * + config.num_decoder_layers * 2 + embedding_flops) * + 3 / 10**12) return attention_tflops, learnable_weight_tflops def calculate_mixed_attention_model_tflops_training_per_device( - config, total_ffn_flops_all_layers, qkv_flops, projection_flops, embedding_flops, attention_pattern_length -): + config, total_ffn_flops_all_layers, qkv_flops, projection_flops, + embedding_flops, attention_pattern_length): """ Calculate training TFLOPs for models with a mixed attention pattern of local and global attention layers, like Gemma3 and GPT-OSS. @@ -311,40 +348,38 @@ def calculate_mixed_attention_model_tflops_training_per_device( # Global causal attention uses a multiplier of 2 (instead of 4 for non-causal) # since we only compute the lower triangular half of the attention matrix. - global_causal_flops_per_layer = ( - 2 * config.per_device_batch_size * config.max_target_length**2 * config.num_query_heads * config.head_dim - ) + global_causal_flops_per_layer = (2 * config.per_device_batch_size * + config.max_target_length**2 * + config.num_query_heads * config.head_dim) # Local sliding window attention directly computes the exact causal interactions # via the formula `(T * W - 0.5 * W^2)`. Therefore, we use the base multiplier of 4. window = min(config.sliding_window_size, config.max_target_length) local_causal_flops_per_layer = ( - 4 - * config.per_device_batch_size - * (config.max_target_length * window - 0.5 * window**2) - * config.num_query_heads - * config.head_dim - ) + 4 * config.per_device_batch_size * + (config.max_target_length * window - 0.5 * window**2) * + config.num_query_heads * config.head_dim) - causal_attention_flops = ( - num_global_layers * global_causal_flops_per_layer + num_local_layers * local_causal_flops_per_layer - ) + causal_attention_flops = (num_global_layers * global_causal_flops_per_layer + + num_local_layers * local_causal_flops_per_layer) # Convert to TFLOPs and multiply by 3 for fwd/bwd pass attention_tflops = causal_attention_flops * 3 / 10**12 total_learnable_flops = total_ffn_flops_all_layers - total_learnable_flops += (qkv_flops + projection_flops) * num_layers + embedding_flops + total_learnable_flops += (qkv_flops + + projection_flops) * num_layers + embedding_flops learnable_weight_tflops = total_learnable_flops * 3 / 10**12 return attention_tflops, learnable_weight_tflops -def calculate_gemma4_tflops_training_per_device( - config, total_ffn_flops_all_layers, embedding_flops, attention_pattern_length -): +def calculate_gemma4_tflops_training_per_device(config, + total_ffn_flops_all_layers, + embedding_flops, + attention_pattern_length): """ Calculate training TFLOPs for Gemma 4. Gemma 4 has specific quirks: @@ -362,70 +397,51 @@ def calculate_gemma4_tflops_training_per_device( # Global causal attention uses a multiplier of 2 (instead of 4 for non-causal) # since we only compute the lower triangular half of the attention matrix. - global_causal_flops_per_layer = ( - 2 * config.per_device_batch_size * config.max_target_length**2 * config.num_query_heads * global_head_dim - ) + global_causal_flops_per_layer = (2 * config.per_device_batch_size * + config.max_target_length**2 * + config.num_query_heads * global_head_dim) # Local sliding window attention directly computes the exact causal interactions # via the formula `(T * W - 0.5 * W^2)`. Therefore, we use the base multiplier of 4. window = min(config.sliding_window_size, config.max_target_length) local_causal_flops_per_layer = ( - 4 - * config.per_device_batch_size - * (config.max_target_length * window - 0.5 * window**2) - * config.num_query_heads - * config.head_dim - ) + 4 * config.per_device_batch_size * + (config.max_target_length * window - 0.5 * window**2) * + config.num_query_heads * config.head_dim) - causal_attention_flops = ( - num_global_layers * global_causal_flops_per_layer + num_local_layers * local_causal_flops_per_layer - ) + causal_attention_flops = (num_global_layers * global_causal_flops_per_layer + + num_local_layers * local_causal_flops_per_layer) # Convert to TFLOPs and multiply by 3 for fwd/bwd pass attention_tflops = causal_attention_flops * 3 / 10**12 global_qkv_flops_per_layer = ( - 2 - * config.per_device_batch_size - * config.max_target_length - * config.emb_dim - * (config.num_query_heads + kv_multiplier * global_num_kv_heads) - * global_head_dim - ) - global_projection_flops_per_layer = ( - 2 - * config.per_device_batch_size - * config.max_target_length - * config.emb_dim - * config.num_query_heads - * global_head_dim - ) + 2 * config.per_device_batch_size * config.max_target_length * + config.emb_dim * + (config.num_query_heads + kv_multiplier * global_num_kv_heads) * + global_head_dim) + global_projection_flops_per_layer = (2 * config.per_device_batch_size * + config.max_target_length * + config.emb_dim * config.num_query_heads * + global_head_dim) # Local layers never share KV projections (kv_multiplier is always 2). local_qkv_flops_per_layer = ( - 2 - * config.per_device_batch_size - * config.max_target_length - * config.emb_dim - * (config.num_query_heads + 2 * config.num_kv_heads) - * config.head_dim - ) - local_projection_flops_per_layer = ( - 2 - * config.per_device_batch_size - * config.max_target_length - * config.emb_dim - * config.num_query_heads - * config.head_dim - ) + 2 * config.per_device_batch_size * config.max_target_length * + config.emb_dim * (config.num_query_heads + 2 * config.num_kv_heads) * + config.head_dim) + local_projection_flops_per_layer = (2 * config.per_device_batch_size * + config.max_target_length * + config.emb_dim * config.num_query_heads * + config.head_dim) total_learnable_flops = total_ffn_flops_all_layers total_learnable_flops += ( - (local_qkv_flops_per_layer + local_projection_flops_per_layer) * num_local_layers - + (global_qkv_flops_per_layer + global_projection_flops_per_layer) * num_global_layers - + embedding_flops - ) + (local_qkv_flops_per_layer + local_projection_flops_per_layer) * + num_local_layers + + (global_qkv_flops_per_layer + global_projection_flops_per_layer) * + num_global_layers + embedding_flops) learnable_weight_tflops = total_learnable_flops * 3 / 10**12 @@ -454,7 +470,8 @@ def calculate_gemma4_small_tflops_training_per_device(config, embedding_flops): num_kv_shared = config.num_kv_shared_layers layer_types = gemma4_small.build_layer_types(num_layers, config.model_name) - first_shared = gemma4_small.first_kv_shared_layer_idx(num_layers, num_kv_shared) + first_shared = gemma4_small.first_kv_shared_layer_idx(num_layers, + num_kv_shared) global_head_dim = config.global_head_dim or config.head_dim global_num_kv_heads = config.global_num_kv_heads or config.num_kv_heads @@ -467,7 +484,8 @@ def attention_flops_for(attn_type): # Causal global attention: factor of 2 (lower-triangular half). return 2 * b * s * s * config.num_query_heads * global_head_dim # Sliding-window: exact closed-form for the causal-and-windowed mask. - return 4 * b * (s * window - 0.5 * window**2) * config.num_query_heads * config.head_dim + return 4 * b * (s * window - + 0.5 * window**2) * config.num_query_heads * config.head_dim def qo_flops_for(attn_type): head_dim = global_head_dim if attn_type == AttentionType.GLOBAL else config.head_dim @@ -508,10 +526,12 @@ def kv_flops_for(attn_type): # One global projection emb_dim -> num_layers * ple_dim. ple_flops += 2 * b * s * config.emb_dim * (num_layers * ple_dim) # Per-layer gate (emb_dim -> ple_dim) and projection (ple_dim -> emb_dim). - ple_flops += num_layers * (2 * b * s * config.emb_dim * ple_dim + 2 * b * s * ple_dim * config.emb_dim) + ple_flops += num_layers * (2 * b * s * config.emb_dim * ple_dim + + 2 * b * s * ple_dim * config.emb_dim) attention_tflops = attention_flops * 3 / 10**12 - learnable_weight_tflops = (projection_flops + ffn_flops + ple_flops + embedding_flops) * 3 / 10**12 + learnable_weight_tflops = (projection_flops + ffn_flops + ple_flops + + embedding_flops) * 3 / 10**12 return attention_tflops, learnable_weight_tflops @@ -542,17 +562,18 @@ def calculate_llama4_attention_tflops(config): num_chunked_layers = num_layers - num_global_layers # FLOPs for a single global attention layer (full attention, non-causal) - global_attention_flops_per_layer = ( - 4 * config.per_device_batch_size * seq_len**2 * config.num_query_heads * config.head_dim - ) + global_attention_flops_per_layer = (4 * config.per_device_batch_size * + seq_len**2 * config.num_query_heads * + config.head_dim) # FLOPs for a single chunked attention layer (non-causal) - chunked_attention_flops_per_layer = _calculate_chunked_attention_flops_per_layer(config, seq_len, chunk_size) + chunked_attention_flops_per_layer = _calculate_chunked_attention_flops_per_layer( + config, seq_len, chunk_size) # Total non-causal attention FLOPs is the sum of all global and all chunked layers - noncausal_attention_flops = (num_global_layers * global_attention_flops_per_layer) + ( - num_chunked_layers * chunked_attention_flops_per_layer - ) + noncausal_attention_flops = ( + num_global_layers * global_attention_flops_per_layer) + ( + num_chunked_layers * chunked_attention_flops_per_layer) # Apply causal mask and convert to TFLOPs (multiply by 3 for fwd/bwd pass) causal_attention_flops = noncausal_attention_flops / 2 @@ -645,47 +666,38 @@ def calculate_mla_tflops_per_device(config): q_flops = 2 * batch_len * config.emb_dim * config.num_query_heads * qk_head_dim_sum else: # calculate query down and up flops - q_flops = ( - 2 - * batch_len - * (config.emb_dim * config.q_lora_rank + config.q_lora_rank * config.num_query_heads * qk_head_dim_sum) - ) + q_flops = (2 * batch_len * + (config.emb_dim * config.q_lora_rank + + config.q_lora_rank * config.num_query_heads * qk_head_dim_sum)) # 2. calculate mla kv projection - kv_flops = ( - 2 - * batch_len - * ( - config.emb_dim * (config.kv_lora_rank + config.qk_rope_head_dim) - + config.kv_lora_rank * config.num_query_heads * (config.qk_nope_head_dim + config.v_head_dim) - ) - ) + kv_flops = (2 * batch_len * (config.emb_dim * + (config.kv_lora_rank + config.qk_rope_head_dim) + + config.kv_lora_rank * config.num_query_heads * + (config.qk_nope_head_dim + config.v_head_dim))) qkv_flops = q_flops + kv_flops # 3. calculate attention if config.use_indexer and config.max_target_length > config.indexer_topk: # get indexer flops - indexer_proj_flops, indexer_scoring_flops = calculate_indexer_tflops_per_device(config) + indexer_proj_flops, indexer_scoring_flops = calculate_indexer_tflops_per_device( + config) qkv_flops += indexer_proj_flops # calculate the proportion of the T x T causal matrix that the Indexer actually explores # this follows the area: (TK - 0.5*K^2) / T^2 (T: max_target_length, K: indexer_topk) - multiplier = calculate_indexer_mask_ratio(config.indexer_topk, config.max_target_length) - attention_flops = ( - 2 - * batch_len - * config.max_target_length - * config.num_query_heads - * (qk_head_dim_sum + config.v_head_dim) - * multiplier - ) + multiplier = calculate_indexer_mask_ratio(config.indexer_topk, + config.max_target_length) + attention_flops = (2 * batch_len * config.max_target_length * + config.num_query_heads * + (qk_head_dim_sum + config.v_head_dim) * multiplier) attention_flops += indexer_scoring_flops else: # standard MLA & max_target_length <= indexer_topk in sparse indexer # in both cases, the indexer is bypassed as the causal mask remains the efficient representation - attention_flops = ( - 2 * batch_len * config.max_target_length * config.num_query_heads * (qk_head_dim_sum + config.v_head_dim) - ) + attention_flops = (2 * batch_len * config.max_target_length * + config.num_query_heads * + (qk_head_dim_sum + config.v_head_dim)) attention_flops = attention_flops / 2 projection_flops = 2 * batch_len * config.emb_dim * config.num_query_heads * config.v_head_dim return qkv_flops, attention_flops, projection_flops @@ -703,9 +715,8 @@ def calculate_ffn_mamtul_tflops_per_device(config, mlp_dim, in_dim=None): """ if in_dim is None: in_dim = config.emb_dim - ffn1_flops = ( - 2 * config.per_device_batch_size * config.max_target_length * mlp_dim * in_dim * len(config.mlp_activations) - ) + ffn1_flops = (2 * config.per_device_batch_size * config.max_target_length * + mlp_dim * in_dim * len(config.mlp_activations)) ffn2_flops = 2 * config.per_device_batch_size * config.max_target_length * mlp_dim * in_dim return ffn1_flops + ffn2_flops @@ -723,12 +734,14 @@ def calculate_routed_and_shared_ffn_tflops_per_device(config): gate_flops = 2 * config.per_device_batch_size * config.max_target_length * config.emb_dim * config.num_experts # Due to the mixed decoder layers, the flops is multiplied by num of layers for both dense and moe num_dense_layers, num_moe_layers = get_dense_moe_layers(config) - dense_ffn_flops = calculate_ffn_mamtul_tflops_per_device(config, config.mlp_dim) * num_dense_layers - shared_experts_flops = ( - calculate_ffn_mamtul_tflops_per_device(config, get_shared_expert_mlp_dim(config)) * config.shared_experts - ) - routed_experts_flops = calculate_ffn_mamtul_tflops_per_device(config, config.moe_mlp_dim) * config.num_experts_per_tok - moe_ffn_flops = (gate_flops + shared_experts_flops + routed_experts_flops) * num_moe_layers + dense_ffn_flops = calculate_ffn_mamtul_tflops_per_device( + config, config.mlp_dim) * num_dense_layers + shared_experts_flops = (calculate_ffn_mamtul_tflops_per_device( + config, get_shared_expert_mlp_dim(config)) * config.shared_experts) + routed_experts_flops = calculate_ffn_mamtul_tflops_per_device( + config, config.moe_mlp_dim) * config.num_experts_per_tok + moe_ffn_flops = (gate_flops + shared_experts_flops + + routed_experts_flops) * num_moe_layers total_ffn_flops = dense_ffn_flops + moe_ffn_flops return total_ffn_flops @@ -743,10 +756,13 @@ def get_dense_moe_layers(config): num_moe_layers = config.num_decoder_layers // config.interleave_moe_layer_step num_dense_layers = config.num_decoder_layers - num_moe_layers return num_dense_layers, num_moe_layers - elif config.decoder_block in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5, DecoderBlockType.DEEPSEEK4): + elif config.decoder_block in (DecoderBlockType.QWEN3_NEXT, + DecoderBlockType.QWEN3_5, + DecoderBlockType.DEEPSEEK4): return 0, config.num_decoder_layers elif config.decoder_block == DecoderBlockType.DEFAULT: - raise ValueError("Unsupported decoder block for dense/MoE layer calculation") + raise ValueError( + "Unsupported decoder block for dense/MoE layer calculation") num_experts = getattr(config, "num_experts", 0) if num_experts > 1: @@ -834,18 +850,22 @@ def calculate_gemma3_vision_layers_tflops_per_device(config): num_patches_h = H // patch_size num_patches_w = W // patch_size seq_len = num_patches_h * num_patches_w # 64*64=4096 - patch_embed_flops = 2 * B * seq_len * (C * patch_size * patch_size) * hidden_dim + patch_embed_flops = 2 * B * seq_len * (C * patch_size * + patch_size) * hidden_dim # 2. gemma3.Encoder: num_layers * gemma3.Encoder1DBlock qkv_flops_per_layer = 3 * (2 * B * seq_len * hidden_dim * hidden_dim) attn_flops_per_layer = 4 * B * seq_len * seq_len * hidden_dim projection_flops_per_layer = 2 * B * seq_len * hidden_dim * hidden_dim # projection after attention multiplication - mlp_flops_per_layer = 2 * (2 * B * seq_len * hidden_dim * intermediate_dim) # two fc layers + mlp_flops_per_layer = 2 * (2 * B * seq_len * hidden_dim * intermediate_dim + ) # two fc layers total_attn_flops = attn_flops_per_layer * num_layers - encoder_flops = (+qkv_flops_per_layer + projection_flops_per_layer + mlp_flops_per_layer) * num_layers + encoder_flops = (+qkv_flops_per_layer + projection_flops_per_layer + + mlp_flops_per_layer) * num_layers # 4. VisionEmbedder - seq_len_after_pooling = (num_patches_h // vision_exit_pooling_window) * (num_patches_w // vision_exit_pooling_window) + seq_len_after_pooling = (num_patches_h // vision_exit_pooling_window) * ( + num_patches_w // vision_exit_pooling_window) vision_embedder_flops = 2 * B * seq_len_after_pooling * hidden_dim * embed_dim # One linear projection # Learnable weights summation @@ -891,16 +911,20 @@ def calculate_llama4_vision_layers_tflops_per_device(config): # 1. Llama4UnfoldConvolution (flops by linear projection) # lax.conv_general_dilated_patches extracts patches through reshaping/indexing without flops # Each patch: C * patch_size * patch_size -> hidden_dim - patch_embed_flops = 2 * B * num_patches * (C * patch_size * patch_size) * hidden_dim + patch_embed_flops = 2 * B * num_patches * (C * patch_size * + patch_size) * hidden_dim # 2. Llama4VisionEncoder: num_layers * (qkv + att_projection + mlp) seq_len = num_patches + 1 # +1 for class token, so 577 - qkv_flops_per_layer = 3 * (2 * B * seq_len * hidden_dim * hidden_dim) # Q, K, V projections + qkv_flops_per_layer = 3 * (2 * B * seq_len * hidden_dim * hidden_dim + ) # Q, K, V projections attn_flops_per_layer = 4 * B * seq_len * seq_len * hidden_dim # Attention scores and weighted sum projection_flops_per_layer = 2 * B * seq_len * hidden_dim * hidden_dim # projection after attention multiplication - mlp_flops_per_layer = 2 * (2 * B * seq_len * hidden_dim * intermediate_dim) # two fc layers + mlp_flops_per_layer = 2 * (2 * B * seq_len * hidden_dim * intermediate_dim + ) # two fc layers total_attn_flops = attn_flops_per_layer * num_layers - vision_encoder_flops = (+qkv_flops_per_layer + projection_flops_per_layer + mlp_flops_per_layer) * num_layers + vision_encoder_flops = (+qkv_flops_per_layer + projection_flops_per_layer + + mlp_flops_per_layer) * num_layers # 3. Llama4VisionPixelShuffleMLP # (B, 144, 5632) -> (B, 144, 4096) -> (B, 144, 4096) @@ -955,7 +979,8 @@ def calculate_engram_tflops(config): num_layers = len(config.engram_layers) # account for both the forward (1x) and backward (2x) passes - learnable_tflops = num_layers * (key_flops + value_flops + conv_flops) * 3 / 1e12 + learnable_tflops = num_layers * (key_flops + value_flops + + conv_flops) * 3 / 1e12 attention_tflops = num_layers * attention_flops * 3 / 1e12 return learnable_tflops, attention_tflops @@ -964,12 +989,10 @@ def calculate_vision_encoder_tflops(config): """Calculate vision encoder TFLOPs per prefill step per device.""" if config.model_name.startswith("gemma3"): mm_total_tflops, mm_learnable_weight_tflops, mm_attention_tflops = calculate_gemma3_vision_layers_tflops_per_device( - config - ) + config) elif config.model_name.startswith("llama4"): mm_total_tflops, mm_learnable_weight_tflops, mm_attention_tflops = calculate_llama4_vision_layers_tflops_per_device( - config - ) + config) else: max_logging.log( f"Vision encoder TFLOPs calculation not implemented for model {config.model_name}, counting as 0 for now." @@ -999,7 +1022,9 @@ def calculate_mhc_flops_per_layer(config): return pre_alpha_flops + pre_contract_flops + post_alpha_flops + post_outer_flops + res_alpha_flops + res_contract_flops -def calculate_deepseek4_tflops_training_per_device(config, total_ffn_flops_all_layers, embedding_flops): +def calculate_deepseek4_tflops_training_per_device(config, + total_ffn_flops_all_layers, + embedding_flops): """Calculates the training TFLOPs per device for DeepSeek-V4. DeepSeek-V4 adopts a multi-track attention topology and Manifold-Constrained @@ -1043,21 +1068,20 @@ def calculate_deepseek4_tflops_training_per_device(config, total_ffn_flops_all_l # 2. Base Attention Block Projections (Executed on all layers) # - Q projection: projects inputs to q_lora_rank, then up-projects to query_heads * head_dim. # FLOPs = 2 * B * S * (D * Q_rank + Q_rank * H * D_h) - wq_flops = ( - 2 - * batch_size - * seq_len - * ((config.emb_dim * config.q_lora_rank) + (config.q_lora_rank * config.num_query_heads * config.head_dim)) - ) + wq_flops = (2 * batch_size * seq_len * + ((config.emb_dim * config.q_lora_rank) + + (config.q_lora_rank * config.num_query_heads * config.head_dim))) # - KV projection: embeds to kv_heads * head_dim. # FLOPs = 2 * B * S * D * (KV_heads * D_h) - wkv_flops = 2 * batch_size * seq_len * config.emb_dim * (config.num_kv_heads * config.head_dim) + wkv_flops = 2 * batch_size * seq_len * config.emb_dim * (config.num_kv_heads * + config.head_dim) # - Grouped Out projection: projects out groups of heads down to o_lora_rank, then up to emb_dim. # FLOPs = 2 * B * S * (H * D_h) * o_lora_rank + 2 * B * S * (o_groups * o_lora_rank) * D grp_feat = (config.num_query_heads * config.head_dim) // config.o_groups - o_flops = (2 * batch_size * seq_len * config.o_groups * grp_feat * config.o_lora_rank) + ( - 2 * batch_size * seq_len * (config.o_groups * config.o_lora_rank) * config.emb_dim - ) + o_flops = (2 * batch_size * seq_len * config.o_groups * grp_feat * + config.o_lora_rank) + (2 * batch_size * seq_len * + (config.o_groups * config.o_lora_rank) * + config.emb_dim) base_attention_weights = wq_flops + wkv_flops + o_flops # 3. Manifold-Constrained Hyper-Connections (mHC) Overheads @@ -1073,36 +1097,31 @@ def calculate_deepseek4_tflops_training_per_device(config, total_ffn_flops_all_l # Uses exact causal sliding window surface area formula to count attention dot products: # FLOPs = 4 * B * (S * W - 0.5 * W^2) * H * D_h prefix_attn = ( - 4 - * batch_size - * (seq_len * sliding_window_size - 0.5 * sliding_window_size**2) - * config.num_query_heads - * config.head_dim - ) + 4 * batch_size * + (seq_len * sliding_window_size - 0.5 * sliding_window_size**2) * + config.num_query_heads * config.head_dim) # - Track 2: HCA Compressed Branch (C=128) # Includes HCA compressor dense projection: D -> D_h hca_pooler = 4 * batch_size * seq_len * config.emb_dim * config.head_dim # Calculates prefix sliding window attention + global cross-attention to the S/128 elements - hca_attn = prefix_attn + 2 * batch_size * seq_len * (seq_len / HCA_RATIO) * config.num_query_heads * config.head_dim + hca_attn = prefix_attn + 2 * batch_size * seq_len * ( + seq_len / HCA_RATIO) * config.num_query_heads * config.head_dim # - Track 3: CSA Sparse Branch (C=4) # - Indexer Projections: Q_rank -> H_idx * D_idx, and D -> H_idx - idx_proj = ( - 2 - * batch_size - * seq_len - * ( - (config.q_lora_rank * config.indexer_n_heads * config.indexer_head_dim) - + (config.emb_dim * config.indexer_n_heads) - ) - ) + idx_proj = (2 * batch_size * seq_len * ( + (config.q_lora_rank * config.indexer_n_heads * config.indexer_head_dim) + + (config.emb_dim * config.indexer_n_heads))) # - Indexer query-block scoring (divided by 2 for causal mask) - idx_score = batch_size * seq_len * (seq_len / CSA_RATIO) * config.indexer_n_heads * config.indexer_head_dim + idx_score = batch_size * seq_len * ( + seq_len / CSA_RATIO) * config.indexer_n_heads * config.indexer_head_dim # - Indexer head reduction (divided by 2 for causal mask) - idx_reduce = batch_size * seq_len * (seq_len / CSA_RATIO) * config.indexer_n_heads + idx_reduce = batch_size * seq_len * (seq_len / + CSA_RATIO) * config.indexer_n_heads # - Indexer key/gate pooler projections: Down-projects D -> 2 * D_idx - idx_kv_gate_pool = 4 * batch_size * seq_len * config.emb_dim * (2 * config.indexer_head_dim) + idx_kv_gate_pool = 4 * batch_size * seq_len * config.emb_dim * ( + 2 * config.indexer_head_dim) csa_indexer = idx_proj + idx_score + idx_reduce + idx_kv_gate_pool # Includes CSA compressor pooler projection: D -> 2 * D_h @@ -1112,22 +1131,19 @@ def calculate_deepseek4_tflops_training_per_device(config, total_ffn_flops_all_l # We reuse the mask multiplier helper for the csa causal ratio: # Ratio = (K/T) - 0.5 * (K/T)^2 csa_mask_ratio = calculate_indexer_mask_ratio(csa_k, seq_len // CSA_RATIO) - csa_sparse_attn = ( - 4 * batch_size * (seq_len * (seq_len / CSA_RATIO)) * config.num_query_heads * config.head_dim * csa_mask_ratio - ) + csa_sparse_attn = (4 * batch_size * (seq_len * (seq_len / CSA_RATIO)) * + config.num_query_heads * config.head_dim * csa_mask_ratio) csa_attn = prefix_attn + csa_sparse_attn + csa_indexer # 5. Pipeline TFLOP Aggregation - total_attn_flops = (num_sliding_layers * prefix_attn) + (num_hca_layers * hca_attn) + (num_csa_layers * csa_attn) + total_attn_flops = (num_sliding_layers * prefix_attn) + ( + num_hca_layers * hca_attn) + (num_csa_layers * csa_attn) total_weight_flops = ( - (num_layers * base_attention_weights) - + (2 * num_layers * mhc_flops_per_layer) # mHC runs twice per layer - + (num_hca_layers * hca_pooler) - + (num_csa_layers * csa_pooler) - + total_ffn_flops_all_layers - + embedding_flops - ) + (num_layers * base_attention_weights) + + (2 * num_layers * mhc_flops_per_layer) # mHC runs twice per layer + + (num_hca_layers * hca_pooler) + (num_csa_layers * csa_pooler) + + total_ffn_flops_all_layers + embedding_flops) # Scale final values by 3x to account for Forward + Backward (2x Forward) training passes return ( @@ -1150,52 +1166,45 @@ def calculate_tflops_training_per_device(config, log=True): DecoderBlockType.GEMMA4, DecoderBlockType.DEEPSEEK4, ): - total_ffn_flops = calculate_routed_and_shared_ffn_tflops_per_device(config) + total_ffn_flops = calculate_routed_and_shared_ffn_tflops_per_device( + config) is_ffn_flops_already_total = True elif config.decoder_block == DecoderBlockType.QWEN3_CUSTOM_MOE: # MoE operates at moe_expert_input_dim (compressed latent), not emb_dim. in_dim = config.moe_expert_input_dim gate_flops = 2 * config.per_device_batch_size * config.max_target_length * in_dim * config.num_experts - total_ffn_flops = ( - gate_flops - + calculate_ffn_mamtul_tflops_per_device(config, config.moe_mlp_dim, in_dim=in_dim) * config.num_experts_per_tok - ) + total_ffn_flops = (gate_flops + calculate_ffn_mamtul_tflops_per_device( + config, config.moe_mlp_dim, in_dim=in_dim) * + config.num_experts_per_tok) else: gate_flops = 2 * config.per_device_batch_size * config.max_target_length * config.emb_dim * config.num_experts - total_ffn_flops = ( - gate_flops + calculate_ffn_mamtul_tflops_per_device(config, config.moe_mlp_dim) * config.num_experts_per_tok - ) + total_ffn_flops = (gate_flops + calculate_ffn_mamtul_tflops_per_device( + config, config.moe_mlp_dim) * config.num_experts_per_tok) else: - total_ffn_flops = calculate_ffn_mamtul_tflops_per_device(config, config.mlp_dim) + total_ffn_flops = calculate_ffn_mamtul_tflops_per_device( + config, config.mlp_dim) - total_ffn_flops_all_layers = ( - total_ffn_flops if is_ffn_flops_already_total else total_ffn_flops * config.num_decoder_layers - ) + total_ffn_flops_all_layers = (total_ffn_flops if is_ffn_flops_already_total + else total_ffn_flops * + config.num_decoder_layers) # Attention flops if config.attention_type == "mla": - qkv_flops, causal_attention_flops, projection_flops = calculate_mla_tflops_per_device(config) + qkv_flops, causal_attention_flops, projection_flops = calculate_mla_tflops_per_device( + config) else: kv_multiplier = 1 if config.share_kv_projections else 2 qkv_flops = ( - 2 - * config.per_device_batch_size - * config.max_target_length - * config.emb_dim - * (config.num_query_heads + kv_multiplier * config.num_kv_heads) - * config.head_dim - ) - noncausal_attention_flops = ( - 4 * config.per_device_batch_size * config.max_target_length**2 * config.num_query_heads * config.head_dim - ) - projection_flops = ( - 2 - * config.per_device_batch_size - * config.max_target_length - * config.emb_dim - * config.num_query_heads - * config.head_dim - ) + 2 * config.per_device_batch_size * config.max_target_length * + config.emb_dim * + (config.num_query_heads + kv_multiplier * config.num_kv_heads) * + config.head_dim) + noncausal_attention_flops = (4 * config.per_device_batch_size * + config.max_target_length**2 * + config.num_query_heads * config.head_dim) + projection_flops = (2 * config.per_device_batch_size * + config.max_target_length * config.emb_dim * + config.num_query_heads * config.head_dim) # Divide attention flops by 2 due to causal mask # References: @@ -1210,60 +1219,63 @@ def calculate_tflops_training_per_device(config, log=True): # Combine flops with number of decoder layers if config.decoder_block == DecoderBlockType.GEMMA2: attention_tflops, learnable_weight_tflops = calculate_gemma2_tflops_training_per_device( - config, total_ffn_flops, qkv_flops, projection_flops, embedding_flops - ) + config, total_ffn_flops, qkv_flops, projection_flops, embedding_flops) elif config.decoder_block == DecoderBlockType.GEMMA3: attention_tflops, learnable_weight_tflops = calculate_mixed_attention_model_tflops_training_per_device( - config, total_ffn_flops_all_layers, qkv_flops, projection_flops, embedding_flops, attention_pattern_length=6 - ) + config, + total_ffn_flops_all_layers, + qkv_flops, + projection_flops, + embedding_flops, + attention_pattern_length=6) elif config.decoder_block == DecoderBlockType.GPT_OSS: attention_tflops, learnable_weight_tflops = calculate_mixed_attention_model_tflops_training_per_device( - config, total_ffn_flops_all_layers, qkv_flops, projection_flops, embedding_flops, attention_pattern_length=2 - ) + config, + total_ffn_flops_all_layers, + qkv_flops, + projection_flops, + embedding_flops, + attention_pattern_length=2) elif config.decoder_block == DecoderBlockType.LLAMA4: # Use the new helper to calculate attention TFLOPs correctly. attention_tflops = calculate_llama4_attention_tflops(config) # The learnable weight calculation remains the same as it correctly handles Llama4's MoE structure. learnable_weight_tflops = ( - (total_ffn_flops_all_layers + (qkv_flops + projection_flops) * config.num_decoder_layers + embedding_flops) - * 3 - / 10**12 - ) + (total_ffn_flops_all_layers + + (qkv_flops + projection_flops) * config.num_decoder_layers + + embedding_flops) * 3 / 10**12) elif config.decoder_block == DecoderBlockType.GEMMA4: attention_tflops, learnable_weight_tflops = calculate_gemma4_tflops_training_per_device( - config, total_ffn_flops_all_layers, embedding_flops, attention_pattern_length=6 - ) + config, + total_ffn_flops_all_layers, + embedding_flops, + attention_pattern_length=6) elif config.decoder_block == DecoderBlockType.GEMMA4_SMALL: # The small path derives its own attention pattern from # gemma4_small.get_attention_pattern(config.model_name) and accounts for # KV sharing, double-wide MLP on shared layers, and the per-layer-embedding # block. - attention_tflops, learnable_weight_tflops = calculate_gemma4_small_tflops_training_per_device(config, embedding_flops) + attention_tflops, learnable_weight_tflops = calculate_gemma4_small_tflops_training_per_device( + config, embedding_flops) elif config.decoder_block == DecoderBlockType.DEEPSEEK4: attention_tflops, learnable_weight_tflops = calculate_deepseek4_tflops_training_per_device( - config, total_ffn_flops_all_layers, embedding_flops - ) + config, total_ffn_flops_all_layers, embedding_flops) elif config.decoder_block == DecoderBlockType.DEEPSEEK: learnable_weight_tflops = ( - (total_ffn_flops_all_layers + (qkv_flops + projection_flops) * config.num_decoder_layers + embedding_flops) - * 3 - / 10**12 - ) + (total_ffn_flops_all_layers + + (qkv_flops + projection_flops) * config.num_decoder_layers + + embedding_flops) * 3 / 10**12) attention_tflops = causal_attention_flops * config.num_decoder_layers * 3 / 10**12 elif config.decoder_block == DecoderBlockType.QWEN3_CUSTOM_MOE: # Attention output projects (num_query_heads * head_dim) -> attention_output_dim, not -> emb_dim. - qwen3_custom_proj_flops = ( - 2 - * config.per_device_batch_size - * config.max_target_length - * config.attention_output_dim - * config.num_query_heads - * config.head_dim - ) + qwen3_custom_proj_flops = (2 * config.per_device_batch_size * + config.max_target_length * + config.attention_output_dim * + config.num_query_heads * config.head_dim) # Each layer has a final up-projection: attention_output_dim -> emb_dim. - layer_up_proj_flops = ( - 2 * config.per_device_batch_size * config.max_target_length * config.attention_output_dim * config.emb_dim - ) + layer_up_proj_flops = (2 * config.per_device_batch_size * + config.max_target_length * + config.attention_output_dim * config.emb_dim) per_layer_flops = qkv_flops + qwen3_custom_proj_flops + layer_up_proj_flops total_weight_flops = total_ffn_flops_all_layers + per_layer_flops * config.num_decoder_layers + embedding_flops learnable_weight_tflops = total_weight_flops * 3 / 10**12 @@ -1272,30 +1284,28 @@ def calculate_tflops_training_per_device(config, log=True): DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5, ): - gdn_weight_flops_per_layer, gdn_attn_flops_per_layer = calculate_gated_delta_net_flops_per_device(config) + gdn_weight_flops_per_layer, gdn_attn_flops_per_layer = calculate_gated_delta_net_flops_per_device( + config) cycle_interval = config.inhomogeneous_layer_cycle_interval num_full_attn_layers = config.num_decoder_layers // cycle_interval num_linear_attn_layers = config.num_decoder_layers - num_full_attn_layers # Weights TFLOPs: - total_weights = ( - total_ffn_flops_all_layers - + embedding_flops - + (qkv_flops + projection_flops) * num_full_attn_layers - + gdn_weight_flops_per_layer * num_linear_attn_layers - ) + total_weights = (total_ffn_flops_all_layers + embedding_flops + + (qkv_flops + projection_flops) * num_full_attn_layers + + gdn_weight_flops_per_layer * num_linear_attn_layers) learnable_weight_tflops = total_weights * 3 / 10**12 # Attention TFLOPs: - total_attn = (causal_attention_flops * num_full_attn_layers) + (gdn_attn_flops_per_layer * num_linear_attn_layers) + total_attn = (causal_attention_flops * num_full_attn_layers) + ( + gdn_attn_flops_per_layer * num_linear_attn_layers) attention_tflops = total_attn * 3 / 10**12 else: # multiply by 3 for both feed forward and back propagation flops learnable_weight_tflops = ( - (total_ffn_flops_all_layers + (qkv_flops + projection_flops) * config.num_decoder_layers + embedding_flops) - * 3 - / 10**12 - ) + (total_ffn_flops_all_layers + + (qkv_flops + projection_flops) * config.num_decoder_layers + + embedding_flops) * 3 / 10**12) attention_tflops = causal_attention_flops * config.num_decoder_layers * 3 / 10**12 # MTP (Multi-Token Prediction) FLOPs Calculation @@ -1314,22 +1324,25 @@ def calculate_tflops_training_per_device(config, log=True): DecoderBlockType.QWEN3_NEXT, DecoderBlockType.GEMMA4, ): - shared_flops = ( - calculate_ffn_mamtul_tflops_per_device(config, get_shared_expert_mlp_dim(config)) * config.shared_experts - ) - routed_flops = calculate_ffn_mamtul_tflops_per_device(config, config.moe_mlp_dim) * config.num_experts_per_tok + shared_flops = (calculate_ffn_mamtul_tflops_per_device( + config, get_shared_expert_mlp_dim(config)) * config.shared_experts) + routed_flops = calculate_ffn_mamtul_tflops_per_device( + config, config.moe_mlp_dim) * config.num_experts_per_tok last_layer_ffn_flops = gate_flops + shared_flops + routed_flops else: - routed_flops = calculate_ffn_mamtul_tflops_per_device(config, config.moe_mlp_dim) * config.num_experts_per_tok + routed_flops = calculate_ffn_mamtul_tflops_per_device( + config, config.moe_mlp_dim) * config.num_experts_per_tok last_layer_ffn_flops = gate_flops + routed_flops else: # For dense architectures, the final layer is a standard dense FFN. - last_layer_ffn_flops = calculate_ffn_mamtul_tflops_per_device(config, config.mlp_dim) + last_layer_ffn_flops = calculate_ffn_mamtul_tflops_per_device( + config, config.mlp_dim) # Calculate total weight FLOPs per MTP module. # Crucially, MTP shares the base model's vocabulary embeddings and final output # projections, so we strictly add only the FFN, QKV, and standard projection FLOPs. - mtp_weight_flops = (last_layer_ffn_flops + qkv_flops + projection_flops) * mtp_num_layers + mtp_weight_flops = (last_layer_ffn_flops + qkv_flops + + projection_flops) * mtp_num_layers # Attention FLOPs scale linearly with the number of MTP modules. mtp_attn_flops = causal_attention_flops * mtp_num_layers @@ -1341,13 +1354,15 @@ def calculate_tflops_training_per_device(config, log=True): # Engram flops if config.engram_layers: - engram_learnable_tflops, engram_attention_tflops = calculate_engram_tflops(config) + engram_learnable_tflops, engram_attention_tflops = calculate_engram_tflops( + config) learnable_weight_tflops += engram_learnable_tflops attention_tflops += engram_attention_tflops if config.use_multimodal: # Add vision layers TFLOPs for multimodal models - mm_total_tflops, mm_learnable_weight_tflops, mm_attention_tflops = calculate_vision_encoder_tflops(config) + mm_total_tflops, mm_learnable_weight_tflops, mm_attention_tflops = calculate_vision_encoder_tflops( + config) if log and mm_total_tflops > 0: print( f"{config.model_name} vision layers per train step:\n", @@ -1383,18 +1398,16 @@ def calculate_tflops_training_per_device(config, log=True): # https://arxiv.org/pdf/2204.02311.pdf Appendix B -def calculate_prefill_tflops_per_device(num_model_parameters, prefill_length, config, log=True): +def calculate_prefill_tflops_per_device(num_model_parameters, + prefill_length, + config, + log=True): """Calculate training TFLOP""" - learnable_weight_tflops = 2 * num_model_parameters * prefill_length / jax.device_count() / 1e12 - noncausal_attention_flops = ( - 4 - * config.num_query_heads - * config.num_decoder_layers - * config.head_dim - * prefill_length**2 - / jax.device_count() - / 1e12 - ) + learnable_weight_tflops = 2 * num_model_parameters * prefill_length / jax.device_count( + ) / 1e12 + noncausal_attention_flops = (4 * config.num_query_heads * + config.num_decoder_layers * config.head_dim * + prefill_length**2 / jax.device_count() / 1e12) causal_attention_tflops = noncausal_attention_flops / 2 # due to causality in attention total_tflops = learnable_weight_tflops + causal_attention_tflops @@ -1455,7 +1468,8 @@ def get_nested_value(dictionary, nested_key, default=None): return current_level -def collect_intermediates_by_suffix(intermediate_outputs, *suffix_keys: str) -> list: +def collect_intermediates_by_suffix(intermediate_outputs, *suffix_keys: + str) -> list: """Collects intermediate leaf values whose dict-key path ends with suffix_keys. Works regardless of model architecture (scanned, scannable blocks, or standard), @@ -1540,7 +1554,11 @@ def _apply_update(path, param): def init_decode_state(apply_fn, params) -> TrainState: """Init train state with null opt state for decode.""" - state = TrainState(step=0, apply_fn=apply_fn, params=params, tx=None, opt_state={}) # type: ignore + state = TrainState(step=0, + apply_fn=apply_fn, + params=params, + tx=None, + opt_state={}) # type: ignore return state @@ -1560,19 +1578,24 @@ def init_initial_state(model, tx, config, is_training, key): """ input_shape = (config.micro_batch_size_to_train_on, config.max_target_length) image_shape = mm_processor.get_dummy_image_shape_for_init( - config.model_name, batch_size=config.micro_batch_size_to_train_on - ) + config.model_name, batch_size=config.micro_batch_size_to_train_on) audio_shape = mm_processor.get_dummy_audio_shape_for_init(config) # Split the master key into independent keys for each RNG collection # Reference: https://flax-linen.readthedocs.io/en/latest/guides/flax_fundamentals/rng_guide.html params_key, dropout_key, aqt_key = jax.random.split(key, 3) model_vars = model.init( - {"params": params_key, "dropout": dropout_key, "aqt": aqt_key}, + { + "params": params_key, + "dropout": dropout_key, + "aqt": aqt_key + }, np.ones(input_shape, dtype=jnp.int32), np.ones(input_shape, dtype=jnp.int32), - encoder_images=np.ones(image_shape, dtype=jnp.int32) if config.use_multimodal else None, - encoder_audios=np.ones(audio_shape, dtype=jnp.float32) if config.use_audio else None, + encoder_images=np.ones(image_shape, dtype=jnp.int32) + if config.use_multimodal else None, + encoder_audios=np.ones(audio_shape, dtype=jnp.float32) + if config.use_audio else None, # nnx_method="no_op", ) if is_training: @@ -1584,26 +1607,30 @@ def get_abstract_param(model, config): """Get abstract model structure (name, shape) without materializing the weights to save memory""" with model.mesh, nn_partitioning.axis_rules(config.logical_axis_rules): key = jax.random.PRNGKey(0) - input_shape = (config.micro_batch_size_to_train_on, config.max_target_length) + input_shape = (config.micro_batch_size_to_train_on, + config.max_target_length) # Cleanly fetch the shapes from the processor using the correct micro_batch_size if config.use_multimodal: image_shape = mm_processor.get_dummy_image_shape_for_init( - config.model_name, batch_size=config.micro_batch_size_to_train_on - ) + config.model_name, batch_size=config.micro_batch_size_to_train_on) else: image_shape = None audio_shape = mm_processor.get_dummy_audio_shape_for_init(config) abstract_vars = jax.eval_shape( model.init, - {"params": key, "dropout": key, "aqt": key}, + { + "params": key, + "dropout": key, + "aqt": key + }, np.ones(input_shape, dtype=jnp.int32), np.ones(input_shape, dtype=jnp.int32), - encoder_images=np.ones(image_shape, dtype=jnp.int32) - if config.use_multimodal - else None, # pyrefly: ignore[no-matching-overload] - encoder_audios=np.ones(audio_shape, dtype=jnp.float32) if config.use_audio else None, + encoder_images=np.ones(image_shape, dtype=jnp.int32) if + config.use_multimodal else None, # pyrefly: ignore[no-matching-overload] + encoder_audios=np.ones(audio_shape, dtype=jnp.float32) + if config.use_audio else None, ) return abstract_vars @@ -1622,14 +1649,15 @@ def setup_decode_state(config, mesh, checkpoint_manager, init_state_fn): """ if not config.load_parameters_path: # generate random params - max_logging.log("No decode checkpoint specified - generating random weights.") + max_logging.log( + "No decode checkpoint specified - generating random weights.") state, state_mesh_annotations, _, _, _ = setup_initial_state( - None, config, mesh, checkpoint_manager, init_state_fn, False - ) + None, config, mesh, checkpoint_manager, init_state_fn, False) else: # Load params from checkpoint max_logging.log(f"Loading decode params from {config.load_parameters_path}") - unboxed_abstract_state, state_mesh_annotations, _ = get_abstract_state(config, mesh, init_state_fn, False) + unboxed_abstract_state, state_mesh_annotations, _ = get_abstract_state( + config, mesh, init_state_fn, False) with nn_partitioning.axis_rules(config.logical_axis_rules): params = checkpointing.load_params_from_path( config.load_parameters_path, @@ -1644,7 +1672,8 @@ def setup_decode_state(config, mesh, checkpoint_manager, init_state_fn): return state, state_mesh_annotations -def setup_training_state(data_iterator, config, mesh, checkpoint_manager, init_state_fn): +def setup_training_state(data_iterator, config, mesh, checkpoint_manager, + init_state_fn): is_training = True return setup_initial_state( data_iterator, @@ -1684,8 +1713,7 @@ def setup_initial_state( """ unboxed_abstract_state, state_mesh_annotations, state_mesh_shardings = get_abstract_state( - config, mesh, init_state_fn, is_training - ) + config, mesh, init_state_fn, is_training) # Initialization with nn_partitioning.axis_rules(config.logical_axis_rules): @@ -1717,7 +1745,8 @@ def setup_initial_state( # checkpoint didn't carry (or a params-only load didn't touch) keeps its real init # value, so restore doesn't depend on knowing exactly what was saved. state = jax.jit( - lambda: nnx.state(init_state_partial()), # Get state only, mapping to out_sharding structure + lambda: nnx.state(init_state_partial( + )), # Get state only, mapping to out_sharding structure in_shardings=None, out_shardings=state_mesh_shardings, )() @@ -1726,7 +1755,8 @@ def setup_initial_state( checkpoint_manager, ( emergency_checkpoint_manager.CheckpointManager, - emergency_replicator_checkpoint_manager.ReplicatorCheckpointManager, + emergency_replicator_checkpoint_manager. + ReplicatorCheckpointManager, ), ) # data_iterator state is updated in place during restore. @@ -1734,13 +1764,7 @@ def setup_initial_state( # didn't carry is still an unmaterialized placeholder. Fill those from the fresh init: a # present leaf comes from the checkpoint, an absent one keeps its init value. overlay = restored if is_emergency else restored["items"] - merged = jax.tree.map( - lambda ckpt, init: init if isinstance(ckpt, jax.ShapeDtypeStruct) else ckpt, - overlay.to_pure_dict(), - state.to_pure_dict(), - is_leaf=lambda x: isinstance(x, jax.ShapeDtypeStruct), - ) - nnx.replace_by_pure_dict(state, merged) + checkpointing._update_nnx_state_from_pure_dict(state, overlay) elif raw_params: # params-only load: overlay the restored weights, keep init for everything else. nnx.update(state.model, raw_params) @@ -1750,7 +1774,8 @@ def setup_initial_state( checkpoint_manager, ( emergency_checkpoint_manager.CheckpointManager, - emergency_replicator_checkpoint_manager.ReplicatorCheckpointManager, + emergency_replicator_checkpoint_manager. + ReplicatorCheckpointManager, ), ): state = restored @@ -1775,7 +1800,8 @@ def _merge_params(p_raw, p_init): return p_init return p_raw - merged_params = jax.tree_util.tree_map(_merge_params, raw_params, state.params) + merged_params = jax.tree_util.tree_map(_merge_params, raw_params, + state.params) state = state.replace(params=merged_params) else: state = state.replace(params=raw_params) @@ -1787,7 +1813,8 @@ def _merge_params(p_raw, p_init): def get_logical_annotations(config, mesh, init_state_fn): init_state_partial = init_state_fn - with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): + with jax.set_mesh(mesh), nn_partitioning.axis_rules( + config.logical_axis_rules): abstract_state = jax.eval_shape(init_state_partial) logical_annotations = nn.get_partition_spec(abstract_state) return logical_annotations @@ -1805,7 +1832,9 @@ def get_abstract_state(config, mesh, init_state_fn, is_training=True): state_logical_annotations = nn.get_partition_spec(abstract_state) - state_mesh_shardings = nn.logical_to_mesh_sharding(state_logical_annotations, mesh, config.logical_axis_rules) + state_mesh_shardings = nn.logical_to_mesh_sharding(state_logical_annotations, + mesh, + config.logical_axis_rules) if is_training and config.shard_optimizer_over_data: # Add data to sharding for optimizer state state_mesh_shardings = state_mesh_shardings.replace( @@ -1813,10 +1842,11 @@ def get_abstract_state(config, mesh, init_state_fn, is_training=True): functools.partial(sharding.add_data_to_sharding, mesh), max_utils.unbox_logicallypartioned(abstract_state).opt_state, state_mesh_shardings.opt_state, - ) - ) + )) if is_training and config.optimizer_memory_host_offload: - opt_state = jax.tree_util.tree_map(lambda x: x.with_memory_kind(kind="pinned_host"), state_mesh_shardings.opt_state) + opt_state = jax.tree_util.tree_map( + lambda x: x.with_memory_kind(kind="pinned_host"), + state_mesh_shardings.opt_state) state_mesh_shardings = state_mesh_shardings.replace(opt_state=opt_state) if is_training and config.parameter_memory_host_offload: assert config.param_scan_axis == 0, "You must set the scan axis 0 to enable parameter offloading." @@ -1828,11 +1858,15 @@ def move(path, x): params = jax.tree_util.tree_map_with_path(move, state_mesh_shardings.params) state_mesh_shardings = state_mesh_shardings.replace(params=params) - abstract_sharded_state = jax.jit(init_state_partial, in_shardings=None, out_shardings=state_mesh_shardings).eval_shape() + abstract_sharded_state = jax.jit( + init_state_partial, in_shardings=None, + out_shardings=state_mesh_shardings).eval_shape() - unboxed_abstract_sharded_state = max_utils.unbox_logicallypartioned(abstract_sharded_state) + unboxed_abstract_sharded_state = max_utils.unbox_logicallypartioned( + abstract_sharded_state) # Initialization - with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): + with jax.set_mesh(mesh), nn_partitioning.axis_rules( + config.logical_axis_rules): state_mesh_annotations = nn.logical_to_mesh(state_logical_annotations) return ( unboxed_abstract_sharded_state, @@ -1841,7 +1875,10 @@ def move(path, x): ) -def get_abstract_state_nnx(config, mesh, nnx_init_trainstate_fn, is_training=True): +def get_abstract_state_nnx(config, + mesh, + nnx_init_trainstate_fn, + is_training=True): """Calculates the abstract sharded state and memory placement for an NNX TrainState. This function performs an abstract trace of the NNX model and optimizer using @@ -1882,14 +1919,16 @@ def get_abstract_state_nnx(config, mesh, nnx_init_trainstate_fn, is_training=Tru # ourselves via nnx_construct_named_sharding, so auto-assignment is not needed here. abs_model = nnx.eval_shape(nnx_init_trainstate_fn) _, abs_var_state = nnx.split(abs_model) - named_sharding_state = sharding.nnx_construct_named_sharding(abs_var_state, mesh) + named_sharding_state = sharding.nnx_construct_named_sharding( + abs_var_state, mesh) abstract_state = jax.tree.map( lambda a, s: jax.ShapeDtypeStruct(a.shape, a.dtype, sharding=s), abs_var_state, named_sharding_state, ) - state_mesh_shardings = maxtext_utils_nnx.nnx_extract_named_sharding(abstract_state) + state_mesh_shardings = maxtext_utils_nnx.nnx_extract_named_sharding( + abstract_state) if is_training and config.shard_optimizer_over_data: # Add data to sharding for optimizer state @@ -1916,8 +1955,10 @@ def get_abstract_state_nnx(config, mesh, nnx_init_trainstate_fn, is_training=Tru ) nnx.update(state_mesh_shardings, state_params) - abstract_sharded_state = maxtext_utils_nnx.set_named_sharding_nnx(abstract_state, state_mesh_shardings) - state_mesh_annotations = maxtext_utils_nnx.get_partition_spec_nnx(state_mesh_shardings) + abstract_sharded_state = maxtext_utils_nnx.set_named_sharding_nnx( + abstract_state, state_mesh_shardings) + state_mesh_annotations = maxtext_utils_nnx.get_partition_spec_nnx( + state_mesh_shardings) return ( abstract_sharded_state, state_mesh_annotations, @@ -1934,12 +1975,15 @@ def init_kv_cache(model, config): config.max_prefill_predict_length, ) image_shape = mm_processor.get_dummy_image_shape_for_init( - config.model_name, batch_size=config.micro_batch_size_to_train_on - ) + config.model_name, batch_size=config.micro_batch_size_to_train_on) audio_shape = mm_processor.get_dummy_audio_shape_for_init(config) model_vars = model.init( - {"params": rng, "dropout": rng, "aqt": rng}, + { + "params": rng, + "dropout": rng, + "aqt": rng + }, jnp.ones(input_shape), jnp.ones(input_shape), encoder_images=jnp.ones(image_shape) if config.use_multimodal else None, @@ -1953,7 +1997,8 @@ def init_kv_cache(model, config): init_kv_cache_partial = functools.partial(init_kv_cache, model, config) abstract_state = jax.eval_shape(init_kv_cache_partial) state_logical_annotations = nn.get_partition_spec(abstract_state) - with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): + with jax.set_mesh(mesh), nn_partitioning.axis_rules( + config.logical_axis_rules): state_mesh_annotations = nn.logical_to_mesh(state_logical_annotations) return state_mesh_annotations @@ -1964,12 +2009,15 @@ def get_kv_cache_annotations(model, config, rng, mesh): def init_kv_cache(model, config): input_shape = (config.micro_batch_size_to_train_on, 1) image_shape = mm_processor.get_dummy_image_shape_for_init( - config.model_name, batch_size=config.micro_batch_size_to_train_on - ) + config.model_name, batch_size=config.micro_batch_size_to_train_on) audio_shape = mm_processor.get_dummy_audio_shape_for_init(config) model_vars = model.init( - {"params": rng, "dropout": rng, "aqt": rng}, + { + "params": rng, + "dropout": rng, + "aqt": rng + }, jnp.ones(input_shape), jnp.ones(input_shape), encoder_images=jnp.ones(image_shape) if config.use_multimodal else None, @@ -1983,7 +2031,8 @@ def init_kv_cache(model, config): init_kv_cache_partial = functools.partial(init_kv_cache, model, config) abstract_state = jax.eval_shape(init_kv_cache_partial) state_logical_annotations = nn.get_partition_spec(abstract_state) - with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules): + with jax.set_mesh(mesh), nn_partitioning.axis_rules( + config.logical_axis_rules): state_mesh_annotations = nn.logical_to_mesh(state_logical_annotations) return state_mesh_annotations @@ -2023,7 +2072,9 @@ def save_quantized_checkpoint_if_configured(config, params): use_zarr3=config.checkpoint_storage_use_zarr3, ) else: - max_logging.log("Skipping saving quantized checkpoint as save_quantized_params_path is null.") + max_logging.log( + "Skipping saving quantized checkpoint as save_quantized_params_path is null." + ) def add_config_to_summary_writer(config, summary_writer): @@ -2051,7 +2102,8 @@ def create_device_mesh(config, devices=None): num_slices = config.num_slices if config.subslice_shape and config.enable_single_controller and num_slices == 1: - max_logging.log(f"Trying to create a subslice with shape: {config.subslice_shape}") + max_logging.log( + f"Trying to create a subslice with shape: {config.subslice_shape}") subslice_shape = tuple(int(x) for x in config.subslice_shape.split(",")) device_coords = [device.coords for device in devices] device_coords_np = np.array(device_coords) @@ -2062,38 +2114,56 @@ def create_device_mesh(config, devices=None): subslice_devices = [] for device in devices: coords = device.coords - if all(min_coords[i] <= coords[i] < min_coords[i] + subslice_shape[i] for i in range(len(subslice_shape))): + if all(min_coords[i] <= coords[i] < min_coords[i] + subslice_shape[i] + for i in range(len(subslice_shape))): subslice_devices.append(device) devices = subslice_devices num_devices = len(devices) - num_slices = 1 if getattr(config, "inference_benchmark_test", False) else num_slices + num_slices = 1 if getattr(config, "inference_benchmark_test", + False) else num_slices num_devices_per_slice = num_devices // num_slices # Find possible unspecified parallelisms ici_parallelism = getattr(config, "ici_parallelism", None) if ici_parallelism is None: ici_map = { - "diloco": getattr(config, "ici_diloco_parallelism", 1), - "data": getattr(config, "ici_data_parallelism", 1), - "stage": getattr(config, "ici_pipeline_parallelism", 1), - "fsdp": getattr(config, "ici_fsdp_parallelism", -1), - "fsdp_transpose": getattr(config, "ici_fsdp_transpose_parallelism", 1), - "sequence": getattr(config, "ici_sequence_parallelism", 1), - "context": getattr(config, "ici_context_parallelism", 1), - "context_autoregressive": getattr(config, "ici_context_autoregressive_parallelism", 1), - "tensor": getattr(config, "ici_tensor_parallelism", 1), - "tensor_sequence": getattr(config, "ici_tensor_sequence_parallelism", 1), - "model": getattr(config, "ici_tensor_parallelism", 1), - "expert": getattr(config, "ici_expert_parallelism", 1), - "autoregressive": getattr(config, "ici_autoregressive_parallelism", 1), - "attn_dp": 1, - "attn_dp_expert": 1, + "diloco": + getattr(config, "ici_diloco_parallelism", 1), + "data": + getattr(config, "ici_data_parallelism", 1), + "stage": + getattr(config, "ici_pipeline_parallelism", 1), + "fsdp": + getattr(config, "ici_fsdp_parallelism", -1), + "fsdp_transpose": + getattr(config, "ici_fsdp_transpose_parallelism", 1), + "sequence": + getattr(config, "ici_sequence_parallelism", 1), + "context": + getattr(config, "ici_context_parallelism", 1), + "context_autoregressive": + getattr(config, "ici_context_autoregressive_parallelism", 1), + "tensor": + getattr(config, "ici_tensor_parallelism", 1), + "tensor_sequence": + getattr(config, "ici_tensor_sequence_parallelism", 1), + "model": + getattr(config, "ici_tensor_parallelism", 1), + "expert": + getattr(config, "ici_expert_parallelism", 1), + "autoregressive": + getattr(config, "ici_autoregressive_parallelism", 1), + "attn_dp": + 1, + "attn_dp_expert": + 1, } ici_parallelism = [ici_map[axis] for axis in config.mesh_axes] else: ici_parallelism = ici_parallelism.copy() - ici_parallelism = max_utils.fill_unspecified_mesh_axes(ici_parallelism, num_devices_per_slice, "ICI") + ici_parallelism = max_utils.fill_unspecified_mesh_axes( + ici_parallelism, num_devices_per_slice, "ICI") allow_split_physical_axes = config.allow_split_physical_axes if config.allow_split_physical_axes else False @@ -2101,29 +2171,47 @@ def create_device_mesh(config, devices=None): dcn_parallelism = getattr(config, "dcn_parallelism", None) if dcn_parallelism is None: dcn_map = { - "diloco": getattr(config, "dcn_diloco_parallelism", 1), - "data": getattr(config, "dcn_data_parallelism", 1), - "stage": getattr(config, "dcn_pipeline_parallelism", 1), - "fsdp": getattr(config, "dcn_fsdp_parallelism", 1), - "fsdp_transpose": getattr(config, "dcn_fsdp_transpose_parallelism", 1), - "sequence": getattr(config, "dcn_sequence_parallelism", 1), - "context": getattr(config, "dcn_context_parallelism", 1), - "context_autoregressive": getattr(config, "dcn_context_autoregressive_parallelism", 1), - "tensor": getattr(config, "dcn_tensor_parallelism", 1), - "tensor_sequence": getattr(config, "dcn_tensor_sequence_parallelism", 1), - "model": getattr(config, "dcn_tensor_parallelism", 1), - "expert": getattr(config, "dcn_expert_parallelism", 1), - "autoregressive": getattr(config, "dcn_autoregressive_parallelism", 1), - "attn_dp": 1, - "attn_dp_expert": 1, + "diloco": + getattr(config, "dcn_diloco_parallelism", 1), + "data": + getattr(config, "dcn_data_parallelism", 1), + "stage": + getattr(config, "dcn_pipeline_parallelism", 1), + "fsdp": + getattr(config, "dcn_fsdp_parallelism", 1), + "fsdp_transpose": + getattr(config, "dcn_fsdp_transpose_parallelism", 1), + "sequence": + getattr(config, "dcn_sequence_parallelism", 1), + "context": + getattr(config, "dcn_context_parallelism", 1), + "context_autoregressive": + getattr(config, "dcn_context_autoregressive_parallelism", 1), + "tensor": + getattr(config, "dcn_tensor_parallelism", 1), + "tensor_sequence": + getattr(config, "dcn_tensor_sequence_parallelism", 1), + "model": + getattr(config, "dcn_tensor_parallelism", 1), + "expert": + getattr(config, "dcn_expert_parallelism", 1), + "autoregressive": + getattr(config, "dcn_autoregressive_parallelism", 1), + "attn_dp": + 1, + "attn_dp_expert": + 1, } dcn_parallelism = [dcn_map[axis] for axis in config.mesh_axes] else: dcn_parallelism = dcn_parallelism.copy() - dcn_parallelism = max_utils.fill_unspecified_mesh_axes(dcn_parallelism, num_slices, "DCN") + dcn_parallelism = max_utils.fill_unspecified_mesh_axes( + dcn_parallelism, num_slices, "DCN") if max_utils.is_valid_custom_mesh(ici_parallelism, config.custom_mesh): - mesh = max_utils.create_custom_device_mesh(ici_parallelism, dcn_parallelism, devices, config.custom_mesh) + mesh = max_utils.create_custom_device_mesh(ici_parallelism, + dcn_parallelism, devices, + config.custom_mesh) else: mesh = mesh_utils.create_hybrid_device_mesh( ici_parallelism, @@ -2184,6 +2272,7 @@ def create_learning_rate_schedule(config): """ def make_cos_schedule(init_lr, final_lr, len_steps): + def schedule(step): pct = step / (len_steps - 1) if len_steps > 1 else 1.0 a = 0.5 * (jnp.cos(jnp.pi * pct) + 1) @@ -2194,14 +2283,17 @@ def schedule(step): lr = config.learning_rate final_lr = lr * config.learning_rate_final_fraction - warmup_steps = int(config.learning_rate_schedule_steps * config.warmup_steps_fraction) + warmup_steps = int(config.learning_rate_schedule_steps * + config.warmup_steps_fraction) constant_zero_steps = config.steps - config.learning_rate_schedule_steps pieces = [] boundaries = [] if warmup_steps > 0: - warmup_schedule = optax.linear_schedule(init_value=0.0, end_value=lr, transition_steps=warmup_steps) + warmup_schedule = optax.linear_schedule(init_value=0.0, + end_value=lr, + transition_steps=warmup_steps) pieces.append(warmup_schedule) boundaries.append(warmup_steps) @@ -2213,7 +2305,8 @@ def schedule(step): boundaries.append(warmup_steps + cos_steps) else: # WSD - decay_steps = int(config.learning_rate_schedule_steps * config.wsd_decay_steps_fraction) + decay_steps = int(config.learning_rate_schedule_steps * + config.wsd_decay_steps_fraction) stable_steps = config.learning_rate_schedule_steps - warmup_steps - decay_steps if stable_steps > 0: @@ -2223,7 +2316,9 @@ def schedule(step): if decay_steps > 0: # Create decay schedule based on wsd_decay_style if config.wsd_decay_style == types.WsdDecayStyle.LINEAR: - decay_schedule = optax.linear_schedule(init_value=lr, end_value=final_lr, transition_steps=decay_steps - 1) + decay_schedule = optax.linear_schedule(init_value=lr, + end_value=final_lr, + transition_steps=decay_steps - 1) else: # COSINE decay_schedule = make_cos_schedule(lr, final_lr, decay_steps) pieces.append(decay_schedule) @@ -2237,7 +2332,10 @@ def schedule(step): return optax.join_schedules(pieces, boundaries) -def print_shardings_params(params, params_sharding, mesh, logical_annotations=None): +def print_shardings_params(params, + params_sharding, + mesh, + logical_annotations=None): """ Print state shardings comparing Logical Definition vs Physical Result. """ @@ -2253,23 +2351,27 @@ def print_shardings_params(params, params_sharding, mesh, logical_annotations=No leaves_sharding, _ = jax.tree_util.tree_flatten_with_path(params_sharding) if logical_annotations is not None: - leaves_logical, _ = jax.tree_util.tree_flatten_with_path(logical_annotations) + leaves_logical, _ = jax.tree_util.tree_flatten_with_path( + logical_annotations) for (path, leaf_val), (_, leaf_sharding), (_, leaf_logical_val) in zip( - leaves_params, leaves_sharding, leaves_logical - ): - path_str = "/".join(str(p.key if hasattr(p, "key") else p.name) for p in path) + leaves_params, leaves_sharding, leaves_logical): + path_str = "/".join( + str(p.key if hasattr(p, "key") else p.name) for p in path) shape = jax.typeof(leaf_val) pspec = sharding.remove_size_one_mesh_axis(leaf_sharding.spec, mesh) pspec_str = str(tuple(pspec)) logical_str = str(leaf_logical_val) - message = ( - f" {path_str}\n" f" Shape: {shape}\n" f" Logical: {logical_str}\n" f" Physical: {pspec_str}" - ) + message = (f" {path_str}\n" + f" Shape: {shape}\n" + f" Logical: {logical_str}\n" + f" Physical: {pspec_str}") max_logging.info(message) else: - for (path, leaf_val), (_, leaf_sharding) in zip(leaves_params, leaves_sharding): - path_str = "/".join(str(p.key if hasattr(p, "key") else p.name) for p in path) + for (path, leaf_val), (_, leaf_sharding) in zip(leaves_params, + leaves_sharding): + path_str = "/".join( + str(p.key if hasattr(p, "key") else p.name) for p in path) shape = jax.typeof(leaf_val) pspec = sharding.remove_size_one_mesh_axis(leaf_sharding.spec, mesh) pspec_str = str(tuple(pspec)) @@ -2298,7 +2400,8 @@ def to_abstract(x): # Convert all input arguments recursively to purely local abstract ShapeDtypeStruct objects # to completely bypass remote Array objects and proxy tracing overhead. abstract_inputs = jax.tree.map(to_abstract, train_step_inputs) - p_train_jaxpr = jax.make_jaxpr(unwrapped_step)(*abstract_inputs) # pyrefly: ignore[no-matching-overload] + p_train_jaxpr = jax.make_jaxpr(unwrapped_step)( + *abstract_inputs) # pyrefly: ignore[no-matching-overload] local_filename = "train_step.jaxpr" local_path = os.path.join(config.dump_jaxpr_local_dir, local_filename) @@ -2316,7 +2419,8 @@ def to_abstract(x): config.dump_jaxpr_local_dir, config.dump_jaxpr_gcs_dir, module_name=local_filename, - delete_local_after=config.dump_jaxpr_delete_local_after, # Keeping local for debugging + delete_local_after=config. + dump_jaxpr_delete_local_after, # Keeping local for debugging all_host_upload=False, # Only upload from lead host (Host 0) ) @@ -2353,16 +2457,24 @@ def prepare_kv_caches_for_scan(kv_caches, scan_length, block_len, stack=False): if not isinstance(kv_caches, list): raise TypeError(f"kv_caches must be a list, got {type(kv_caches)}") - grouped = [tuple(kv_caches[i * block_len : (i + 1) * block_len]) for i in range(scan_length)] + grouped = [ + tuple(kv_caches[i * block_len:(i + 1) * block_len]) + for i in range(scan_length) + ] if stack: # Stack the list of tuples into a tuple of stacked PyTrees along a new leading axis - return jax.tree_util.tree_map(lambda *args: jnp.stack(args, axis=0), *grouped) + return jax.tree_util.tree_map(lambda *args: jnp.stack(args, axis=0), + *grouped) else: return grouped -def update_kv_caches_after_scan(kv_caches, returned_kv_cache, scan_length, block_len, stacked=False): +def update_kv_caches_after_scan(kv_caches, + returned_kv_cache, + scan_length, + block_len, + stacked=False): """Updates the original flat list of KV caches from the scanned outputs.""" if kv_caches is None or returned_kv_cache is None: return @@ -2372,7 +2484,10 @@ def update_kv_caches_after_scan(kv_caches, returned_kv_cache, scan_length, block if stacked: # Unstack the returned tuple of stacked PyTrees back into a list of tuples - unstacked = [jax.tree_util.tree_map(lambda x, i=i: x[i], returned_kv_cache) for i in range(scan_length)] + unstacked = [ + jax.tree_util.tree_map(lambda x, i=i: x[i], returned_kv_cache) + for i in range(scan_length) + ] for i in range(scan_length): start_idx = i * block_len for offset, updated_item in enumerate(unstacked[i]): diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py index bacdb631d2..e44b99529b 100644 --- a/src/maxtext/utils/sharding.py +++ b/src/maxtext/utils/sharding.py @@ -45,9 +45,13 @@ def get_input_data_sharding(config, mesh, rules=None): if rules is None: rules = config.logical_axis_rules if config.enable_diloco: - data_sharding = create_sharding(mesh, ["diloco"] + config.input_data_sharding_logical_axes, rules=rules) + data_sharding = create_sharding(mesh, ["diloco"] + + config.input_data_sharding_logical_axes, + rules=rules) else: - data_sharding = create_sharding(mesh, config.input_data_sharding_logical_axes, rules=rules) + data_sharding = create_sharding(mesh, + config.input_data_sharding_logical_axes, + rules=rules) return data_sharding @@ -66,7 +70,10 @@ def _get_sharding_desc(inputs, extra_stack_level): if frame is not None: callers_local_vars = frame.f_locals.items() - x = [var_name for var_name, var_val in callers_local_vars if var_val is inputs] + x = [ + var_name for var_name, var_val in callers_local_vars + if var_val is inputs + ] if len(x) > 0: caller_path_full = inspect.stack()[1 + extra_stack_level].filename # Use pathlib.Path to easily extract just the filename from the full path. @@ -75,9 +82,13 @@ def _get_sharding_desc(inputs, extra_stack_level): return "Unknown" -def maybe_shard_with_name( - inputs, named_sharding, shard_mode, debug_sharding=False, extra_stack_level=0, sharding_desc="", logical_axes=None -): +def maybe_shard_with_name(inputs, + named_sharding, + shard_mode, + debug_sharding=False, + extra_stack_level=0, + sharding_desc="", + logical_axes=None): """ In auto shardmode, this function hints inputs follow given named_sharding. In explicit shardmode, this function enforces inputs following named_sharding. @@ -86,9 +97,8 @@ def maybe_shard_with_name( """ if inputs is None: return None - if ( - debug_sharding and isinstance(inputs, Tracer) and isinstance(named_sharding, NamedSharding) - ): # only print pspec for JitTracer + if (debug_sharding and isinstance(inputs, Tracer) and isinstance( + named_sharding, NamedSharding)): # only print pspec for JitTracer if not sharding_desc: sharding_desc = _get_sharding_desc(inputs, extra_stack_level + 1) @@ -97,30 +107,37 @@ def maybe_shard_with_name( elif isinstance(logical_axes, list): logical_axes = tuple(logical_axes) - pspec = remove_size_one_mesh_axis(getattr(named_sharding, "spec"), getattr(named_sharding, "mesh")) - log_key = (sharding_desc, str(jax.typeof(inputs)), tuple(pspec), extra_stack_level) + pspec = remove_size_one_mesh_axis(getattr(named_sharding, "spec"), + getattr(named_sharding, "mesh")) + log_key = (sharding_desc, str(jax.typeof(inputs)), tuple(pspec), + extra_stack_level) if log_key not in _LOGGED_ACTIVATION_SHARDINGS: - max_logging.info(f"{sharding_desc} Logical: {log_key[1]:.<60} {logical_axes}.", stacklevel=3 + extra_stack_level) - max_logging.info(f"{sharding_desc} Physical: {log_key[1]:.<60} {log_key[2]}.", stacklevel=3 + extra_stack_level) + max_logging.info( + f"{sharding_desc} Logical: {log_key[1]:.<60} {logical_axes}.", + stacklevel=3 + extra_stack_level) + max_logging.info( + f"{sharding_desc} Physical: {log_key[1]:.<60} {log_key[2]}.", + stacklevel=3 + extra_stack_level) _LOGGED_ACTIVATION_SHARDINGS.add(log_key) - _ACTIVATION_SHARDINGS_DUMP.append( - { - f"{sharding_desc}: {log_key[1]}": { - "logic_axes": f"{logical_axes}", - "PartitionSpec": f"P{log_key[2]}", - } + _ACTIVATION_SHARDINGS_DUMP.append({ + f"{sharding_desc}: {log_key[1]}": { + "logic_axes": f"{logical_axes}", + "PartitionSpec": f"P{log_key[2]}", } - ) + }) if shard_mode == ShardMode.EXPLICIT: return reshard(inputs, named_sharding) else: return jax.lax.with_sharding_constraint(inputs, named_sharding) -def maybe_shard_with_pspec( - inputs, pspec: jax.sharding.PartitionSpec | None, mesh, shard_mode, debug_sharding=False, extra_stack_level=0 -): +def maybe_shard_with_pspec(inputs, + pspec: jax.sharding.PartitionSpec | None, + mesh, + shard_mode, + debug_sharding=False, + extra_stack_level=0): if pspec is None: return None sharding = NamedSharding(mesh, pspec) @@ -133,9 +150,14 @@ def maybe_shard_with_pspec( ) -def maybe_shard_with_logical( - inputs, logical_axes, mesh, shard_mode, rules=None, debug_sharding=False, extra_stack_level=0, sharding_desc="" -): +def maybe_shard_with_logical(inputs, + logical_axes, + mesh, + shard_mode, + rules=None, + debug_sharding=False, + extra_stack_level=0, + sharding_desc=""): """ A wrapper of maybe_shard_with_name when logical axes are inputs sharding_desc is description of inputs of upper layer(s) of caller (with the form of /). @@ -181,7 +203,8 @@ def remove_size_one_mesh_axis(spec, mesh): return P(*new_spec, unreduced=spec.unreduced, reduced=spec.reduced) -def get_nnx_var_named_sharding_with_scan_axis(v: nnx.Variable, mesh) -> nnx.Variable: +def get_nnx_var_named_sharding_with_scan_axis(v: nnx.Variable, + mesh) -> nnx.Variable: """Compute NamedSharding for an NNX variable, correctly handling the scan axis.""" val = v.get_value() if not hasattr(val, "shape"): @@ -193,7 +216,8 @@ def get_nnx_var_named_sharding_with_scan_axis(v: nnx.Variable, mesh) -> nnx.Vari return v.replace(jax.tree.map(lambda _: replicated, val)) return v metadata = v.get_metadata() - out_sharding = metadata.get("out_sharding") or metadata.get("sharding_names") or metadata.get("sharding") + out_sharding = metadata.get("out_sharding") or metadata.get( + "sharding_names") or metadata.get("sharding") if not out_sharding: pspec = P() else: @@ -201,7 +225,8 @@ def get_nnx_var_named_sharding_with_scan_axis(v: nnx.Variable, mesh) -> nnx.Vari if nnx.PARTITION_NAME in metadata: partition_name = metadata[nnx.PARTITION_NAME] scan_axis = metadata.get("param_scan_axis", 0) - out_sharding = [out_sharding] if isinstance(out_sharding, str) else list(out_sharding) + out_sharding = [out_sharding] if isinstance(out_sharding, + str) else list(out_sharding) if partition_name not in out_sharding: out_sharding.insert(scan_axis, partition_name) out_sharding = tuple(out_sharding) @@ -210,13 +235,16 @@ def get_nnx_var_named_sharding_with_scan_axis(v: nnx.Variable, mesh) -> nnx.Vari local_rules = metadata.get("sharding_rules", ()) if context_rules or local_rules: local_rules_list = list(local_rules) if local_rules is not None else [] - context_rules_list = list(context_rules) if context_rules is not None else [] + context_rules_list = list( + context_rules) if context_rules is not None else [] rules = local_rules_list + context_rules_list pspec = logical_to_mesh_axes(out_sharding, mesh, rules=rules) else: pspec = P(*out_sharding) if mesh is not None: pspec = remove_size_one_mesh_axis(pspec, mesh) + if mesh is None: + return v.replace(pspec) return v.replace(NamedSharding(mesh, pspec)) @@ -283,7 +311,8 @@ def logical_to_mesh_sharding(tree, mesh, rules=None): def create_sharding(mesh, logical_names, rules=None): """Create NamedSharding with given logical names.""" - return NamedSharding(mesh, logical_to_mesh_axes(logical_names, mesh, rules=rules)) + return NamedSharding(mesh, + logical_to_mesh_axes(logical_names, mesh, rules=rules)) def get_mesh_axes_used_by_tensor_spec(tensor_sharding_spec): @@ -305,10 +334,9 @@ def get_mesh_axes_used_by_tensor_spec(tensor_sharding_spec): """ # Flatten the sharding spec, as it can contain nested iterables (e.g., ('data', 'mdl')). tensor_sharding_spec = sum( - [ - [axis] if isinstance(axis, str) else list(axis) if isinstance(axis, Iterable) else [] - for axis in tensor_sharding_spec - ], + [[axis] if isinstance(axis, str) else + list(axis) if isinstance(axis, Iterable) else [] + for axis in tensor_sharding_spec], [], ) return tensor_sharding_spec @@ -344,7 +372,10 @@ def _get_nontrival_mesh_axes(mesh): # Filter the target axes to find those that exist in the current mesh # and have a size greater than 1, meaning they are actually used for sharding. - return {axis for axis in target_sharding_axes_config if axis in mesh.axis_names and mesh.shape[axis] > 1} + return { + axis for axis in target_sharding_axes_config + if axis in mesh.axis_names and mesh.shape[axis] > 1 + } def _analyze_sharding(params, mesh, valid_target_mesh_axes): @@ -368,48 +399,50 @@ def _analyze_sharding(params, mesh, valid_target_mesh_axes): dictionary contains details about a tensor that is not sharded on any of the target axes. """ unsharded_params_total_size = 0 # Initialize a counter for the size of unsharded parameters. - problematic_tensors_details = [] # Initialize a list to store details of problematic tensors. + problematic_tensors_details = [ + ] # Initialize a list to store details of problematic tensors. # Get a flattened list of all parameters (leaves) in the PyTree, along with their paths. all_params_leaves = jax.tree_util.tree_leaves_with_path(params) for path, p_leaf in all_params_leaves: # Iterate over each parameter leaf - param_name_str = jax.tree_util.keystr(path) # Convert the tree path to a readable string + param_name_str = jax.tree_util.keystr( + path) # Convert the tree path to a readable string # Check that sharding and spec exist and are valid sharding = getattr(p_leaf, "sharding", None) spec = getattr(sharding, "spec", None) assert sharding is not None and spec is not None and isinstance(spec, P), ( f"Parameter '{param_name_str}' is missing a valid '.sharding.spec'." - "Expected 'p_leaf.sharding.spec' to be a non-null 'partitionspec'." - ) + "Expected 'p_leaf.sharding.spec' to be a non-null 'partitionspec'.") current_sharding_spec = p_leaf.sharding.spec # Extract the current tensor's sharding spec # Identify axes used for sharding mesh_axes_used = get_mesh_axes_used_by_tensor_spec(current_sharding_spec) # Check if the parameter is sharded on all the valid target axes. - is_sharded_on_all_target_axis = all(axis in mesh_axes_used for axis in valid_target_mesh_axes) + is_sharded_on_all_target_axis = all( + axis in mesh_axes_used for axis in valid_target_mesh_axes) # If the parameter is not sharded on all of the target axes, it's considered "problematic." if not is_sharded_on_all_target_axis: unsharded_params_total_size += p_leaf.size # Add to total unsharded parameter size unsharded_axes = set(valid_target_mesh_axes) - set(mesh_axes_used) # Add detailed info to list of problematic tensors - problematic_tensors_details.append( - { - "name": param_name_str, # Tensor name - "size": p_leaf.size, # tensor size - "shape": p_leaf.shape, # tensor shape - "spec": str(current_sharding_spec), # Tensor sharding spec as string - "available_axes": sorted(list(valid_target_mesh_axes)), # Axes that could be used for sharding - "unsharded_axes": sorted(list(unsharded_axes)), # Unsharded axes - } - ) + problematic_tensors_details.append({ + "name": param_name_str, # Tensor name + "size": p_leaf.size, # tensor size + "shape": p_leaf.shape, # tensor shape + "spec": str(current_sharding_spec), # Tensor sharding spec as string + "available_axes": sorted(list(valid_target_mesh_axes) + ), # Axes that could be used for sharding + "unsharded_axes": sorted(list(unsharded_axes)), # Unsharded axes + }) # Return the total size of unsharded parameters and the list of problematic tensors. return unsharded_params_total_size, problematic_tensors_details # Return results -def _raise_if_unsharded_exceeds_tolerance(unsharded_size, total_size, tolerance, problematic_tensors_details): +def _raise_if_unsharded_exceeds_tolerance(unsharded_size, total_size, tolerance, + problematic_tensors_details): """ Raises an AssertionError if the percentage of unsharded parameters exceeds the given tolerance. @@ -440,20 +473,20 @@ def _raise_if_unsharded_exceeds_tolerance(unsharded_size, total_size, tolerance, # Begin constructing the error message. error_msg_lines = [ - f"Unsharded parameter percentage ({unsharded_param_perc:.2%})" f"exceeds tolerance ({tolerance:.2%})." + f"Unsharded parameter percentage ({unsharded_param_perc:.2%})" + f"exceeds tolerance ({tolerance:.2%})." ] # Add a header explaining the issue. error_msg_lines.append( "The following large tensors are replicated (unsharded) but could be sharded on at " - "least one of the available axes:" - ) + "least one of the available axes:") # Add details for the top 5 largest problematic tensors. - for detail in problematic_tensors_details[:5]: # Show top 5 largest problematic tensors + for detail in problematic_tensors_details[: + 5]: # Show top 5 largest problematic tensors error_msg_lines.append( f" - Name: {detail['name']}(Size: {detail['size']}, Shape: {detail['spec']}, Spec: {detail['spec']}) " f" is unsharded on axis: {detail['unsharded_axes']}" - f" could be sharded on: {detail['available_axes']}" - ) + f" could be sharded on: {detail['available_axes']}") # Raise the assertion error with the combined, formatted message. raise AssertionError("\n".join(error_msg_lines)) @@ -484,13 +517,14 @@ def assert_params_sufficiently_sharded(params, mesh, tolerance): # Analyze the parameters to find the total size of unsharded parameters # and get details on which tensors are problematic. - unsharded_params_total_size, problematic_tensors_details = _analyze_sharding(params, mesh, valid_target_mesh_axes) + unsharded_params_total_size, problematic_tensors_details = _analyze_sharding( + params, mesh, valid_target_mesh_axes) # Check if the amount of unsharded parameters is within the tolerance and # raise an exception if it is not. - _raise_if_unsharded_exceeds_tolerance( - unsharded_params_total_size, total_num_params, tolerance, problematic_tensors_details - ) + _raise_if_unsharded_exceeds_tolerance(unsharded_params_total_size, + total_num_params, tolerance, + problematic_tensors_details) def add_data_to_sharding(mesh, path, aval, sharding): @@ -514,11 +548,15 @@ def add_data_to_sharding(mesh, path, aval, sharding): AssertionError: If sharding is not NamedSharding or shape cannot be sharded """ if not isinstance(sharding, jax.sharding.NamedSharding): - raise AssertionError(f"Expected NamedSharding, found {sharding} of {type(sharding)=} at {jax.tree_util.keystr(path)}") + raise AssertionError( + f"Expected NamedSharding, found {sharding} of {type(sharding)=} at {jax.tree_util.keystr(path)}" + ) try: sharded_shape = sharding.shard_shape(aval.shape) except Exception as e: - raise AssertionError(f"Could not shard {jax.tree_util.keystr(path)} of shape={aval.shape} with {sharding=}") from e + raise AssertionError( + f"Could not shard {jax.tree_util.keystr(path)} of shape={aval.shape} with {sharding=}" + ) from e pspec = sharding.spec if "data" in jax.tree.leaves(pspec): @@ -531,9 +569,12 @@ def add_data_to_sharding(mesh, path, aval, sharding): if isinstance(partition, str): partition = (partition,) - if size % mesh.shape["data"] == 0 and (partition is None or "tensor" not in partition): - added_component = ("data",) + partition # pyrefly: ignore[unsupported-operation] - new_pspec = jax.sharding.PartitionSpec(*(pspec[:idx] + (added_component,) + pspec[idx + 1 :])) + if size % mesh.shape["data"] == 0 and (partition is None or + "tensor" not in partition): + added_component = ( + "data",) + partition # pyrefly: ignore[unsupported-operation] + new_pspec = jax.sharding.PartitionSpec( + *(pspec[:idx] + (added_component,) + pspec[idx + 1:])) new_sharding = jax.sharding.NamedSharding(sharding.mesh, new_pspec) return new_sharding return sharding @@ -558,30 +599,32 @@ def maybe_update_params_sharding_with_opt(config, state_mesh_shardings): (unchanged if shard_optimizer_over_data is False) """ if config.pure_nnx: - return maybe_update_params_sharding_with_opt_nnx(config, state_mesh_shardings) + return maybe_update_params_sharding_with_opt_nnx(config, + state_mesh_shardings) prev_params_shardings = state_mesh_shardings.params if config.shard_optimizer_over_data: if isinstance(state_mesh_shardings.opt_state, optax.ScaleByAdamState): sharded_fp32_params = state_mesh_shardings.opt_state.mu elif isinstance(state_mesh_shardings.opt_state, tuple) and isinstance( - state_mesh_shardings.opt_state[0], optax.ScaleByAdamState - ): + state_mesh_shardings.opt_state[0], optax.ScaleByAdamState): sharded_fp32_params = state_mesh_shardings.opt_state[0].mu else: - raise NotImplementedError(f"Could not find optimizer state shardings from {type(state_mesh_shardings.opt_state)}") + raise NotImplementedError( + f"Could not find optimizer state shardings from {type(state_mesh_shardings.opt_state)}" + ) if "params" not in sharded_fp32_params.keys(): # When quantization=fp8 is enabled the sharded_fp32_params # are not wrapped in `params`. Here we wrap them back. sharded_fp32_params = {"params": sharded_fp32_params} - state_mesh_shardings = state_mesh_shardings.replace( - params=dict(prev_params_shardings, **sharded_fp32_params) - ) # pyrefly: ignore[bad-unpacking] + state_mesh_shardings = state_mesh_shardings.replace(params=dict( + prev_params_shardings, + **sharded_fp32_params)) # pyrefly: ignore[bad-unpacking] return prev_params_shardings, state_mesh_shardings def maybe_update_params_sharding_with_opt_nnx( - config: pyconfig.HyperParameters, state_mesh_shardings: nnx.State -) -> tuple[nnx.State, nnx.State]: + config: pyconfig.HyperParameters, + state_mesh_shardings: nnx.State) -> tuple[nnx.State, nnx.State]: """ NNX version of parameter sharding update. Updates parameter sharding configuration when optimizer state sharding is enabled. @@ -610,9 +653,13 @@ def _extract_param_only(state): structure as nnx.split(model, nnx.Param, ...)[1], enabling jax.tree.map to work correctly between ga_params (Param-only) and params_shardings. """ + param_type = nnx.LoRAParam if getattr(getattr(config, "lora", None), + "enable_lora", False) else nnx.Param result = {} for k, v in state.items(): - if isinstance(v, nnx.Param): + if isinstance(v, (NamedSharding, jax.sharding.Sharding)): + result[k] = v + elif isinstance(v, param_type): result[k] = v elif isinstance(v, nnx.Variable): pass # skip non-Param variables (RngKey, RngCount, OptVariable, etc.) @@ -663,24 +710,40 @@ def find_adam_mu(obj): sharded_fp32_params = find_adam_mu(opt_state) if sharded_fp32_params is None: actual_type = type(state_mesh_shardings.optimizer.get("opt_state", "None")) - raise NotImplementedError(f"Could not find Adam optimizer state in: {actual_type}") + raise NotImplementedError( + f"Could not find Adam optimizer state in: {actual_type}") # Update model parameter sharding to match the mu (first moment) sharding. # This ensures parameter sharding is consistent with the Zero-1 distributed layout. # Build a path → new_PS lookup from sharded_fp32_params (mu), then update model_shardings # at those paths while preserving rngs and any other non-Param variables. mu_leaves_with_paths = list( - jax.tree_util.tree_leaves_with_path(sharded_fp32_params, is_leaf=lambda x: isinstance(x, nnx.Variable)) - ) - mu_lookup = {path: mu_var.get_value() for path, mu_var in mu_leaves_with_paths} + jax.tree_util.tree_leaves_with_path( + sharded_fp32_params, + is_leaf=lambda x: isinstance(x, (nnx.Variable, NamedSharding, jax. + sharding.Sharding)), + )) + mu_lookup = { + path: (mu_var.get_value() if isinstance(mu_var, nnx.Variable) else mu_var) + for path, mu_var in mu_leaves_with_paths + } def _update_model_var(path, var): if path in mu_lookup: - return var.replace(mu_lookup[path]) + if isinstance(var, nnx.Variable): + val = var.get_value() + if isinstance(val, jax.ShapeDtypeStruct): + return var.replace(value=jax.ShapeDtypeStruct( + val.shape, val.dtype, sharding=mu_lookup[path])) + return var.replace(mu_lookup[path]) + return mu_lookup[path] return var new_model_shardings = jax.tree_util.tree_map_with_path( - _update_model_var, model_shardings, is_leaf=lambda x: isinstance(x, nnx.Variable) + _update_model_var, + model_shardings, + is_leaf=lambda x: isinstance(x, (nnx.Variable, NamedSharding, jax.sharding + .Sharding)), ) # Use jax.tree_util.tree_map (identity) to create a new nnx.State via JAX's unflatten # mechanism (not the nnx.State constructor). This is critical because: @@ -689,13 +752,17 @@ def _update_model_var(path, var): # nested module states as plain dicts). JAX's unflatten preserves the original types. # 2. copy.deepcopy fails because NamedSharding contains non-picklable jaxlib.Device objects. # Direct __setattr__ assignment stores new_model_shardings as-is (no type conversion). - updated_state = jax.tree_util.tree_map(lambda x: x, state_mesh_shardings, is_leaf=lambda x: isinstance(x, nnx.Variable)) + updated_state = jax.tree_util.tree_map( + lambda x: x, + state_mesh_shardings, + is_leaf=lambda x: isinstance(x, nnx.Variable)) updated_state.model = new_model_shardings return prev_params_shardings, updated_state -def build_zero1_input_state_mesh_shardings(config, state_mesh_shardings, params_shardings): +def build_zero1_input_state_mesh_shardings(config, state_mesh_shardings, + params_shardings): """Build the train-step input shardings under shard_optimizer_over_data (Zero-1). Model params on input use the original pre-Zero-1 sharding (params_shardings), @@ -742,7 +809,8 @@ def logical_axis_rules_pp_act_as_dp(logical_rules): new_physical_axes = tuple(axis for axis in physical_axes if axis != "stage") if "data" in new_physical_axes: data_idx = new_physical_axes.index("data") - new_physical_axes = new_physical_axes[0:data_idx] + ("stage",) + new_physical_axes[data_idx:] + new_physical_axes = new_physical_axes[0:data_idx] + ( + "stage",) + new_physical_axes[data_idx:] new_rules.append((key, new_physical_axes)) return tuple(new_rules) @@ -769,7 +837,8 @@ def get_formatted_sharding_annotations(params, mesh=None): # If a mesh object is provided, add its details to the report header. if mesh: - annotation_lines.append(f"Mesh axes: {mesh.axis_names}, Mesh shape: {mesh.shape}") + annotation_lines.append( + f"Mesh axes: {mesh.axis_names}, Mesh shape: {mesh.shape}") annotation_lines.append("-" * 30) # Get a flattened list of all parameters (leaves) and their corresponding paths in the PyTree. @@ -806,7 +875,9 @@ def get_formatted_sharding_annotations(params, mesh=None): sharding_desc = "No .sharding attribute found" # Append the formatted details for the current parameter to our list of lines. - annotation_lines.append(f" - Param: {param_name_str}\n" f" Shape: {shape_str}\n" f" Sharding: {sharding_desc}") + annotation_lines.append(f" - Param: {param_name_str}\n" + f" Shape: {shape_str}\n" + f" Sharding: {sharding_desc}") # Join all the collected lines into a single string, separated by newlines. return "\n".join(annotation_lines) @@ -835,7 +906,8 @@ def _remove_fsdp_from_partition_spec(named_sharding): else: raise ValueError(f"Unsupported_axis_type: {type(axis)}") # Return a new sharding object with the modified spec. - return jax.sharding.NamedSharding(named_sharding.mesh, jax.sharding.PartitionSpec(*new_spec)) + return jax.sharding.NamedSharding(named_sharding.mesh, + jax.sharding.PartitionSpec(*new_spec)) return named_sharding return jax.tree.map(_remove_fsdp_from_partition_spec, sharding_tree) @@ -892,13 +964,16 @@ def get_physical_spec_no_fsdp(full_logical, mesh, logical_axis_rules): """ # Convert the high-level logical spec to a physical one using default rules. - physical = logical_to_mesh_sharding(full_logical, mesh=mesh, rules=logical_axis_rules) + physical = logical_to_mesh_sharding(full_logical, + mesh=mesh, + rules=logical_axis_rules) # Apply the function to remove the FSDP sharding, defining our target layout. physical_no_fsdp = remove_fsdp_sharding(physical) return physical_no_fsdp -def all_gather_over_fsdp(variables, sharding_info, mesh, logical_axis_rules, shard_mode): +def all_gather_over_fsdp(variables, sharding_info, mesh, logical_axis_rules, + shard_mode): """Performs an all-gather on FSDP-sharded variables via a sharding constraint. This function triggers an all-gather operation on the model's parameters. It does so by applying a sharding constraint that specifies a fully @@ -920,7 +995,10 @@ def all_gather_over_fsdp(variables, sharding_info, mesh, logical_axis_rules, sha in the weights being fully replicated on all devices in the 'fsdp' mesh. """ # Get the target physical layout (weights fully replicated). - physical_constraint_no_fsdp = get_physical_spec_no_fsdp(sharding_info, mesh, logical_axis_rules) + physical_constraint_no_fsdp = get_physical_spec_no_fsdp( + sharding_info, mesh, logical_axis_rules) # Apply the constraint to the model's current variables. This tells JAX to # gather the weights into this layout. - return maybe_shard_with_name(variables, physical_constraint_no_fsdp, shard_mode=shard_mode) + return maybe_shard_with_name(variables, + physical_constraint_no_fsdp, + shard_mode=shard_mode) diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py index bfc21de4fe..92c890bf11 100644 --- a/src/maxtext/utils/train_utils.py +++ b/src/maxtext/utils/train_utils.py @@ -32,6 +32,7 @@ from maxtext.common.goodput import GoodputEvent, maybe_record_goodput from maxtext.optimizers import optimizers from maxtext.trainers.diloco import diloco +from maxtext.utils import lora_utils from maxtext.utils import max_logging from maxtext.utils import max_utils from maxtext.utils import maxtext_utils @@ -234,17 +235,27 @@ def setup_train_loop(config, recorder, devices=None): context_parallel_size = mesh.shape.get(config.context_sharding, 1) if config.pure_nnx: # Create abstract NNX model. - _create_model_partial, model = model_creation_utils.create_nnx_abstract_model(config, mesh, devices) + _create_model_partial, _ = model_creation_utils.create_nnx_abstract_model(config, mesh, devices) + model = _create_model_partial() + if getattr(getattr(config, "lora", None), "enable_lora", False): + model = lora_utils.apply_lora_to_model(model, mesh, config) else: model = model_creation_utils.from_config(config, devices) + learning_rate_schedule, tx = create_training_optimizer(config, model) if config.pure_nnx: # For NNX, the train state is wrapped in the TrainStateNNX module. def create_train_state_fn(): - model = _create_model_partial() - optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param) - return train_state_nnx.TrainStateNNX(model, optimizer) + m = _create_model_partial() + if getattr(getattr(config, "lora", None), "enable_lora", False): + m = lora_utils.apply_lora_to_model(m, mesh, config) + param_types = (getattr(nnx, "LoRAParam", nnx.Param),) + else: + param_types = (nnx.Param,) + nnx.pop(m, nnx.Intermediate) + optimizer = nnx.Optimizer(m, tx, wrt=param_types) + return train_state_nnx.TrainStateNNX(m, optimizer) init_state_fn = create_train_state_fn else: @@ -311,9 +322,13 @@ def create_train_state_fn(): # under jax.set_mesh(mesh) and rejects any logical name missing from # logical_axis_rules (e.g. concat_embed on the MTP kernel). Tracing shapes # without a mesh skips sharding resolution, so it avoids the crash. - state_graphdef = nnx.graphdef(nnx.eval_shape(init_state_fn)) - _, state_params, _ = nnx.split(state.model, nnx.Param, ...) - _, state_mesh_shardings_params, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...) + state_graphdef = nnx.graphdef(init_state_fn()) + if isinstance(state, nnx.State): + state_params = state["model"] + state_mesh_shardings_params = state_mesh_shardings["model"] + else: + state_params = state.model + state_mesh_shardings_params = state_mesh_shardings.model else: state_params = state.params state_mesh_shardings_params = state_mesh_shardings.params @@ -363,7 +378,10 @@ def create_train_state_fn(): train_state = state model = state_graphdef # pyrefly: ignore[unbound-name] else: - train_state = nnx.merge(state_graphdef, state) # pyrefly: ignore[unbound-name] + if isinstance(state, nnx.State): + train_state = nnx.merge(state_graphdef, state) + else: + train_state = state model = train_state.model else: train_state = state diff --git a/tests/integration/peft_integration_test.py b/tests/integration/peft_integration_test.py new file mode 100644 index 0000000000..98f1eaf911 --- /dev/null +++ b/tests/integration/peft_integration_test.py @@ -0,0 +1,654 @@ +# 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. + +# pylint: disable=import-outside-toplevel, protected-access +"""Integration tests for Flax NNX PEFT, LoRA, and QLoRA.""" + +import os +import sys +import unittest +import tempfile +import json +import shutil +import pytest +import subprocess + +from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT +from tests.utils.test_helpers import get_test_config_path + + +def _get_jax_backend(): + """Determine JAX backend by running a quick subprocess to avoid locking TPU in the parent process.""" + try: + result = subprocess.run( + [sys.executable, "-c", "import jax; print(jax.default_backend())"], + capture_output=True, + text=True, + check=True) + return result.stdout.strip() + except Exception: # pylint: disable=broad-exception-caught + return "cpu" + + +def _get_jax_device_count(): + """Determine JAX device count by running a quick subprocess.""" + try: + result = subprocess.run( + [sys.executable, "-c", "import jax; print(jax.device_count())"], + capture_output=True, + text=True, + check=True) + return int(result.stdout.strip()) + except Exception: # pylint: disable=broad-exception-caught + return 1 + + +def _get_common_args(temp_dir, model_name, steps=1, overrides=None): + """Build common arguments for downscaled integration runs.""" + tokenizer_path = os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizers", + "tokenizer.llama2") + + backend = _get_jax_backend() + num_devices = _get_jax_device_count() + + hardware = "cpu" + if backend == "tpu": + hardware = "tpu" + elif backend == "gpu": + hardware = "gpu" + + args = [ + None, + get_test_config_path(), + f"run_name=test_peft_{model_name}", + f"model_name={model_name}", + f"steps={steps}", + f"base_output_directory={temp_dir}", + "dataset_type=synthetic", + "per_device_batch_size=1", + "max_target_length=32", + f"tokenizer_path={tokenizer_path}", + "enable_goodput_recording=False", + "enable_checkpoint_cloud_logger=False", + "monitor_goodput=False", + "sharding_tolerance=1.0", # Bypass sharding checks for downscaled test models + f"hardware={hardware}", + "skip_jax_distributed_system=True", + "ici_fsdp_parallelism=1", + f"ici_data_parallelism={num_devices}", + ] + if overrides: + args.extend(overrides) + return args + + +def _run_train_subprocess(args): + """Run the pre-training trainer inside an isolated Python subprocess.""" + cmd = [sys.executable, "-m", "maxtext.trainers.pre_train.train", args[1] + ] + args[2:] + env = os.environ.copy() + env["PYTHONPATH"] = os.path.abspath("src") + + backend = _get_jax_backend() + + if backend not in ("tpu", "gpu"): + env["JAX_PLATFORMS"] = "cpu" + + subprocess.run(cmd, env=env, check=True) + + +def _run_sft_subprocess(args): + """Run the SFT native trainer inside an isolated Python subprocess.""" + cmd = [ + sys.executable, "-m", "maxtext.trainers.post_train.sft.train_sft_native", + args[1] + ] + args[2:] + env = os.environ.copy() + env["PYTHONPATH"] = os.path.abspath("src") + + backend = _get_jax_backend() + + if backend not in ("tpu", "gpu"): + env["JAX_PLATFORMS"] = "cpu" + + subprocess.run(cmd, env=env, check=True) + + +@pytest.mark.integration_test +class PEFTIntegrationTest(unittest.TestCase): + """Integration checks for pre-training and native SFT with LoRA & QLoRA.""" + + def setUp(self): + self.temp_dirs = [] + + def tearDown(self): + for temp_dir in self.temp_dirs: + if os.path.exists(temp_dir): + shutil.rmtree(temp_dir) + + def _create_temp_dir(self): + temp_dir = tempfile.mkdtemp() + self.temp_dirs.append(temp_dir) + return temp_dir + + def test_native_sft_gemma4_lora(self): + """Test native SFT with Gemma4-e2b and LoRA enabled.""" + temp_dir = self._create_temp_dir() + overrides = [ + "override_model_config=True", + "scan_layers=False", + "weight_dtype=bfloat16", + "dtype=bfloat16", + "lora.enable_lora=True", + "lora.lora_rank=4", + "lora.lora_alpha=8.0", + "base_num_decoder_layers=2", + "base_emb_dim=256", + "base_mlp_dim=512", + "num_kv_shared_layers=1", + ] + args = _get_common_args(temp_dir, + "gemma4-e2b", + steps=2, + overrides=overrides) + _run_sft_subprocess(args) + + def test_native_sft_gemma4_qlora(self): + """Test native SFT with Gemma4-e2b and QLoRA enabled.""" + temp_dir = self._create_temp_dir() + overrides = [ + "override_model_config=True", + "scan_layers=False", + "weight_dtype=bfloat16", + "dtype=bfloat16", + "lora.enable_lora=True", + "lora.lora_weight_qtype=int8", + "lora.lora_tile_size=32", + "lora.lora_rank=4", + "lora.lora_alpha=8.0", + "base_num_decoder_layers=2", + "base_emb_dim=256", + "base_mlp_dim=512", + "num_kv_shared_layers=1", + ] + args = _get_common_args(temp_dir, + "gemma4-e2b", + steps=2, + overrides=overrides) + _run_sft_subprocess(args) + + def test_native_sft_gemma4_lora_checkpoint_save_and_restore(self): + """Test native SFT with Gemma4-e2b, LoRA enabled, checkpoint save and subsequent restore.""" + temp_dir = self._create_temp_dir() + metrics_file_saved = os.path.join(temp_dir, "saved_metrics_gemma4.txt") + metrics_file_restored = os.path.join(temp_dir, + "restored_metrics_gemma4.txt") + + overrides_base = [ + "override_model_config=True", + "scan_layers=False", + "weight_dtype=bfloat16", + "dtype=bfloat16", + "lora.enable_lora=True", + "lora.lora_rank=4", + "lora.lora_alpha=8.0", + "base_num_decoder_layers=2", + "base_emb_dim=256", + "base_mlp_dim=512", + "num_kv_shared_layers=1", + ] + + # 1. Step 1 run with checkpointing enabled + args_save = _get_common_args( + temp_dir, + "gemma4-e2b", + steps=1, + overrides=overrides_base + [ + "enable_checkpointing=True", + "checkpoint_period=1", + "async_checkpointing=False", + f"metrics_file={metrics_file_saved}", + ], + ) + _run_sft_subprocess(args_save) + + checkpoint_dir = os.path.join(temp_dir, "test_peft_gemma4-e2b", + "checkpoints", "0", "items") + self.assertTrue(os.path.exists(checkpoint_dir), + f"Checkpoint directory {checkpoint_dir} was not created.") + + # 2. Step 2 run restoring from step 1 checkpoint + args_restore = _get_common_args( + temp_dir, + "gemma4-e2b", + steps=2, + overrides=overrides_base + [ + "enable_checkpointing=True", + "checkpoint_period=1", + "async_checkpointing=False", + f"load_parameters_path={checkpoint_dir}", + f"metrics_file={metrics_file_restored}", + ], + ) + _run_sft_subprocess(args_restore) + + self.assertTrue(os.path.exists(metrics_file_saved)) + self.assertTrue(os.path.exists(metrics_file_restored)) + + with open(metrics_file_saved, "r", encoding="utf8") as f: + saved_metrics = [json.loads(line) for line in f.readlines()] + with open(metrics_file_restored, "r", encoding="utf8") as f: + restored_metrics = [json.loads(line) for line in f.readlines()] + + self.assertEqual(len(saved_metrics), 1) + self.assertEqual(len(restored_metrics), 1) + self.assertEqual(saved_metrics[0]["step"], 0.0) + self.assertEqual(restored_metrics[0]["step"], 1.0) + + def test_native_sft_qwen3_lora(self): + """Test native SFT with Qwen3-0.6B and LoRA enabled.""" + temp_dir = self._create_temp_dir() + overrides = [ + "override_model_config=True", + "scan_layers=True", + "weight_dtype=bfloat16", + "dtype=bfloat16", + "lora.enable_lora=True", + "lora.lora_rank=4", + "lora.lora_alpha=8.0", + "base_num_decoder_layers=2", + "base_emb_dim=256", + "base_mlp_dim=512", + ] + args = _get_common_args(temp_dir, + "qwen3-0.6b", + steps=2, + overrides=overrides) + _run_sft_subprocess(args) + + def test_native_sft_qwen3_qlora(self): + """Test native SFT with Qwen3-0.6B and QLoRA enabled.""" + temp_dir = self._create_temp_dir() + overrides = [ + "override_model_config=True", + "scan_layers=True", + "weight_dtype=bfloat16", + "dtype=bfloat16", + "lora.enable_lora=True", + "lora.lora_weight_qtype=int8", + "lora.lora_tile_size=32", + "lora.lora_rank=4", + "lora.lora_alpha=8.0", + "base_num_decoder_layers=2", + "base_emb_dim=256", + "base_mlp_dim=512", + ] + args = _get_common_args(temp_dir, + "qwen3-0.6b", + steps=2, + overrides=overrides) + _run_sft_subprocess(args) + + def test_qlora_applied_verification(self): + """Explicitly verify that QLoRA applies LoRAParam adapters and quantizes base weights.""" + from flax import nnx + import jax + from maxtext.utils import lora_utils, model_creation_utils + from maxtext.configs import pyconfig + + cfg = pyconfig.initialize_pydantic( + [sys.argv[0], get_test_config_path()], + run_name="test_qlora_verify", + model_name="llama2-7b", + override_model_config=True, + scan_layers=False, + base_num_decoder_layers=1, + base_emb_dim=128, + base_mlp_dim=256, + lora={ + "enable_lora": True, + "lora_rank": 4, + "lora_alpha": 8.0, + "lora_weight_qtype": "int8", + "lora_tile_size": 32, + }, + ) + + model, mesh = model_creation_utils.from_pretrained( + cfg, mesh=None, model_mode=model_creation_utils.MODEL_MODE_TRAIN) + lora_model = lora_utils.apply_lora_to_model(model, mesh, cfg) + + # 1. Assert model is identified as LoRA-enabled + self.assertTrue(lora_utils.is_lora_enabled(lora_model)) + + # 2. Assert LoRAParam adapter variables exist + _, state = nnx.split(lora_model) + lora_params = nnx.filter_state(state, nnx.LoRAParam) + self.assertGreater(len(jax.tree_util.tree_leaves(lora_params)), 0) + + # 3. Assert QLoRA quantization config is active + self.assertEqual(cfg.lora.lora_weight_qtype, "int8") + + def test_pretrain_lora_checkpoint_save_and_restore(self): + """Test pre-training with LoRA, checkpoint saving, and subsequent restoration.""" + temp_dir = self._create_temp_dir() + metrics_file_saved = os.path.join(temp_dir, "saved_metrics.txt") + metrics_file_restored = os.path.join(temp_dir, "restored_metrics.txt") + + # Common model scaling/configs + overrides_base = [ + "scan_layers=True", + "attention=dot_product", + "weight_dtype=bfloat16", + "dtype=bfloat16", + "lora.enable_lora=True", + "lora.lora_rank=4", + "lora.lora_alpha=8.0", + "base_num_decoder_layers=2", + "base_emb_dim=128", + "base_mlp_dim=256", + "ici_fsdp_parallelism=1", + ] + + # 1. First run: train 1 step with checkpointing enabled + args_save = _get_common_args( + temp_dir, + "default", + steps=1, + overrides=overrides_base + [ + "enable_checkpointing=True", + "checkpoint_period=1", + "async_checkpointing=False", + f"metrics_file={metrics_file_saved}", + ], + ) + _run_train_subprocess(args_save) + + # Verify checkpoint got written (MaxText writes to base_output_directory/run_name/checkpoints/0/items) + checkpoint_dir = os.path.join(temp_dir, "test_peft_default", "checkpoints", + "0", "items") + self.assertTrue(os.path.exists(checkpoint_dir), + f"Checkpoint directory {checkpoint_dir} was not created.") + + # 2. Second run: restore from the checkpoint directory and train 1 step (total steps = 2) + args_restore = _get_common_args( + temp_dir, + "default", + steps=2, + overrides=overrides_base + [ + "enable_checkpointing=True", + "checkpoint_period=1", + "async_checkpointing=False", + f"load_parameters_path={checkpoint_dir}", + f"metrics_file={metrics_file_restored}", + ], + ) + _run_train_subprocess(args_restore) + + # 3. Verify losses are logged and correct + self.assertTrue(os.path.exists(metrics_file_saved)) + self.assertTrue(os.path.exists(metrics_file_restored)) + + with open(metrics_file_saved, "r", encoding="utf8") as f: + saved_metrics = [json.loads(line) for line in f.readlines()] + with open(metrics_file_restored, "r", encoding="utf8") as f: + restored_metrics = [json.loads(line) for line in f.readlines()] + + print("SAVED METRICS:", saved_metrics) + print("RESTORED METRICS:", restored_metrics) + + self.assertEqual(len(saved_metrics), 1, + "First run should have exactly 1 step.") + self.assertEqual(len(restored_metrics), 1, + "Second run should have exactly 1 step.") + + saved_step0_loss = saved_metrics[0]["learning/loss"] + restored_step1_loss = restored_metrics[0]["learning/loss"] + + print(f"Saved run Step 0 Loss: {saved_step0_loss}") + print(f"Restored run Step 1 Loss: {restored_step1_loss}") + + # Verify that the restored run resumed from step 0 (logged step 1) + self.assertEqual(saved_metrics[0]["step"], 0.0) + self.assertEqual(restored_metrics[0]["step"], 1.0) + + def test_pretrain_qlora_checkpoint_save_and_restore(self): + """Test pre-training with QLoRA, checkpoint saving, and subsequent restoration.""" + temp_dir = self._create_temp_dir() + metrics_file_saved = os.path.join(temp_dir, "saved_metrics_qlora.txt") + metrics_file_restored = os.path.join(temp_dir, "restored_metrics_qlora.txt") + + # Common model scaling/configs + overrides_base = [ + "scan_layers=True", + "attention=dot_product", + "weight_dtype=bfloat16", + "dtype=bfloat16", + "lora.enable_lora=True", + "lora.lora_weight_qtype=int8", + "lora.lora_tile_size=32", + "lora.lora_rank=4", + "lora.lora_alpha=8.0", + "base_num_decoder_layers=2", + "base_emb_dim=128", + "base_mlp_dim=256", + "ici_fsdp_parallelism=1", + ] + + # 1. First run: train 1 step with checkpointing enabled + args_save = _get_common_args( + temp_dir, + "default", + steps=1, + overrides=overrides_base + [ + "enable_checkpointing=True", + "checkpoint_period=1", + "async_checkpointing=False", + f"metrics_file={metrics_file_saved}", + ], + ) + _run_train_subprocess(args_save) + + # Verify checkpoint got written (MaxText writes to base_output_directory/run_name/checkpoints/0/items) + checkpoint_dir = os.path.join(temp_dir, "test_peft_default", "checkpoints", + "0", "items") + self.assertTrue(os.path.exists(checkpoint_dir), + f"Checkpoint directory {checkpoint_dir} was not created.") + + # 2. Second run: restore from the checkpoint directory and train 1 step (total steps = 2) + args_restore = _get_common_args( + temp_dir, + "default", + steps=2, + overrides=overrides_base + [ + "enable_checkpointing=True", + "checkpoint_period=1", + "async_checkpointing=False", + f"load_parameters_path={checkpoint_dir}", + f"metrics_file={metrics_file_restored}", + ], + ) + _run_train_subprocess(args_restore) + + # 3. Verify losses are logged and correct + self.assertTrue(os.path.exists(metrics_file_saved)) + self.assertTrue(os.path.exists(metrics_file_restored)) + + with open(metrics_file_saved, "r", encoding="utf8") as f: + saved_metrics = [json.loads(line) for line in f.readlines()] + with open(metrics_file_restored, "r", encoding="utf8") as f: + restored_metrics = [json.loads(line) for line in f.readlines()] + + print("SAVED METRICS (QLoRA):", saved_metrics) + print("RESTORED METRICS (QLoRA):", restored_metrics) + + self.assertEqual(len(saved_metrics), 1, + "First run should have exactly 1 step.") + self.assertEqual(len(restored_metrics), 1, + "Second run should have exactly 1 step.") + + saved_step0_loss = saved_metrics[0]["learning/loss"] + restored_step1_loss = restored_metrics[0]["learning/loss"] + + print(f"Saved run Step 0 Loss (QLoRA): {saved_step0_loss}") + print(f"Restored run Step 1 Loss (QLoRA): {restored_step1_loss}") + + # Verify that the restored run resumed from step 0 (logged step 1) + self.assertEqual(saved_metrics[0]["step"], 0.0) + self.assertEqual(restored_metrics[0]["step"], 1.0) + + def test_lora_optimizer_wrt_filtering(self): + """Verify optimizer wrt filtering targets nnx.LoRAParam and base params remain frozen.""" + from flax import nnx + import jax + from maxtext.utils import lora_utils, model_creation_utils + from maxtext.configs import pyconfig + + cfg = pyconfig.initialize_pydantic( + [sys.argv[0], get_test_config_path()], + run_name="test_wrt_filter", + model_name="llama2-7b", + override_model_config=True, + scan_layers=False, + base_num_decoder_layers=1, + base_emb_dim=128, + base_mlp_dim=256, + lora={ + "enable_lora": True, + "lora_rank": 4, + "lora_alpha": 8.0, + }, + ) + + model, mesh = model_creation_utils.from_pretrained( + cfg, mesh=None, model_mode=model_creation_utils.MODEL_MODE_TRAIN) + lora_model = lora_utils.apply_lora_to_model(model, mesh, cfg) + + # Filter parameters by type + _, state = nnx.split(lora_model) + lora_params = nnx.filter_state(state, nnx.LoRAParam) + base_params = nnx.filter_state(state, nnx.Param) + + # Verify LoRAParam exists and base_params does not include LoRAParam + self.assertGreater(len(jax.tree_util.tree_leaves(lora_params)), 0) + self.assertGreater(len(jax.tree_util.tree_leaves(base_params)), 0) + + # Verify that wrt filtering works as expected for optax/nnx optimizer + wrt_filter = nnx.LoRAParam + filtered_state = nnx.filter_state(state, wrt_filter) + self.assertEqual(len(jax.tree_util.tree_leaves(filtered_state)), + len(jax.tree_util.tree_leaves(lora_params))) + + def test_standalone_lora_adapter_restoration(self): + """Test loading base parameters and separate LoRA adapter weights.""" + from flax import nnx + from maxtext.utils import lora_utils, model_creation_utils + from maxtext.common import checkpointing + from maxtext.configs import pyconfig + + cfg = pyconfig.initialize_pydantic( + [sys.argv[0], get_test_config_path()], + run_name="test_adapter_restore", + model_name="llama2-7b", + override_model_config=True, + scan_layers=False, + base_num_decoder_layers=1, + base_emb_dim=128, + base_mlp_dim=256, + lora={ + "enable_lora": True, + "lora_rank": 4, + "lora_alpha": 8.0, + }, + ) + + model, mesh = model_creation_utils.from_pretrained( + cfg, mesh=None, model_mode=model_creation_utils.MODEL_MODE_TRAIN) + lora_model = lora_utils.apply_lora_to_model(model, mesh, cfg) + + # Extract target lora parameter structure + target_state = nnx.state(lora_model, (nnx.Param, nnx.LoRAParam)) + has_pure = hasattr(target_state, "to_pure_dict") + pure_target = target_state.to_pure_dict( + ) if has_pure else checkpointing._tree_to_dict(target_state) + + self.assertIn("decoder", pure_target) + + def test_pretrain_gemma4_scanned_blocks_lora_resume(self): + """Test Gemma4 pre-training with LoRA, checkpointing, and resume.""" + temp_dir = self._create_temp_dir() + metrics_file_saved = os.path.join(temp_dir, "saved_metrics_gemma4_scan.txt") + metrics_file_restored = os.path.join(temp_dir, + "restored_metrics_gemma4_scan.txt") + + overrides_base = [ + "override_model_config=True", + "scan_layers=False", + "weight_dtype=bfloat16", + "dtype=bfloat16", + "lora.enable_lora=True", + "lora.lora_rank=4", + "lora.lora_alpha=8.0", + "base_num_decoder_layers=2", + "base_emb_dim=256", + "base_mlp_dim=512", + "num_kv_shared_layers=1", + ] + + # 1. Step 1: Save checkpoint + args_save = _get_common_args( + temp_dir, + "gemma4-e2b", + steps=1, + overrides=overrides_base + [ + "enable_checkpointing=True", + "checkpoint_period=1", + "async_checkpointing=False", + f"metrics_file={metrics_file_saved}", + ], + ) + _run_train_subprocess(args_save) + + checkpoint_dir = os.path.join(temp_dir, "test_peft_gemma4-e2b", + "checkpoints", "0", "items") + self.assertTrue(os.path.exists(checkpoint_dir), + f"Checkpoint directory {checkpoint_dir} was not created.") + + # 2. Step 2: Resume training from step 1 checkpoint + args_restore = _get_common_args( + temp_dir, + "gemma4-e2b", + steps=2, + overrides=overrides_base + [ + "enable_checkpointing=True", + "checkpoint_period=1", + "async_checkpointing=False", + f"load_parameters_path={checkpoint_dir}", + f"metrics_file={metrics_file_restored}", + ], + ) + _run_train_subprocess(args_restore) + + self.assertTrue(os.path.exists(metrics_file_saved)) + self.assertTrue(os.path.exists(metrics_file_restored)) + + with open(metrics_file_saved, "r", encoding="utf8") as f: + saved_metrics = [json.loads(line) for line in f.readlines()] + with open(metrics_file_restored, "r", encoding="utf8") as f: + restored_metrics = [json.loads(line) for line in f.readlines()] + + self.assertEqual(len(saved_metrics), 1) + self.assertEqual(len(restored_metrics), 1) + self.assertEqual(saved_metrics[0]["step"], 0.0) + self.assertEqual(restored_metrics[0]["step"], 1.0) diff --git a/tests/integration/setup_train_loop_nnx_test.py b/tests/integration/setup_train_loop_nnx_test.py index f446c8f9df..c2e1c09022 100644 --- a/tests/integration/setup_train_loop_nnx_test.py +++ b/tests/integration/setup_train_loop_nnx_test.py @@ -28,6 +28,7 @@ from flax import nnx import jax +from jax.sharding import NamedSharding from maxtext.common import train_state_nnx from maxtext.configs import pyconfig from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT @@ -118,15 +119,11 @@ def test_pure_nnx_setup_param_only_split_matches_model(self): *_, state_mesh_shardings, model, _, _, _, _, _, _, train_state = setup_train_loop(config, recorder=None) _, params, _ = nnx.split(train_state.model, nnx.Param, ...) - _, params_shardings, _ = nnx.split(state_mesh_shardings.model, nnx.Param, ...) - - # Same key-set after nnx.split — this is what setup_train_loop relies on at - # train_utils.py:281-282 to pair state_params with state_mesh_shardings_params. - self.assertEqual( - jax.tree_util.tree_structure(params), - jax.tree_util.tree_structure(params_shardings), - ) - self.assertGreater(len(jax.tree.leaves(params)), 0) + params_flat = dict(nnx.to_flat_state(params)) + full_shardings_flat = dict(nnx.to_flat_state(state_mesh_shardings.model)) + params_shardings_flat = {k: full_shardings_flat[k] for k in params_flat if k in full_shardings_flat} + self.assertEqual(set(params_flat.keys()), set(params_shardings_flat.keys())) + self.assertGreater(len(params_flat), 0) del model diff --git a/tests/post_training/unit/lora_utils_test.py b/tests/post_training/unit/lora_utils_test.py index 01f39feed5..50e5af9308 100644 --- a/tests/post_training/unit/lora_utils_test.py +++ b/tests/post_training/unit/lora_utils_test.py @@ -12,7 +12,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +# pylint: disable=import-outside-toplevel, protected-access """Tests for Qwix LoRA utils in lora_utils.py""" + import re import sys import tempfile @@ -22,6 +24,7 @@ from etils import epath import jax import jax.numpy as jnp +import numpy as np import optax import pytest from flax import nnx @@ -154,8 +157,10 @@ def test_verify_lora_parameters_success(self): ] with ( - mock.patch("maxtext.utils.lora_utils.nnx.iter_graph", return_value=mock_graph_entries), - mock.patch("maxtext.utils.lora_utils.is_lora_enabled", return_value=True), + mock.patch("maxtext.utils.lora_utils.nnx.iter_graph", + return_value=mock_graph_entries), + mock.patch("maxtext.utils.lora_utils.is_lora_enabled", + return_value=True), mock.patch("maxtext.utils.max_logging.log") as mock_log, ): lora_utils._verify_lora_parameters(mock_model, mock_config) @@ -163,7 +168,8 @@ def test_verify_lora_parameters_success(self): # Should log the successful match pattern summary log_calls = [call[0][0] for call in mock_log.call_args_list] self.assertTrue(any("successfully matched" in msg for msg in log_calls)) - self.assertTrue(any("Sample matched submodules" in msg for msg in log_calls)) + self.assertTrue( + any("Sample matched submodules" in msg for msg in log_calls)) def test_verify_lora_parameters_not_enabled_no_match(self): """Test verification fails with ValueError when no modules match at all.""" @@ -174,7 +180,8 @@ def test_verify_lora_parameters_not_enabled_no_match(self): with ( mock.patch("flax.nnx.iter_modules", return_value=[]), - mock.patch("maxtext.utils.lora_utils.is_lora_enabled", return_value=False), + mock.patch("maxtext.utils.lora_utils.is_lora_enabled", + return_value=False), ): with self.assertRaisesRegex(ValueError, "no LoRA parameters found"): lora_utils._verify_lora_parameters(mock_model, mock_config) @@ -182,7 +189,8 @@ def test_verify_lora_parameters_not_enabled_no_match(self): def test_apply_lora_to_model_disabled(self): """Test applying LoRA when it is disabled in config.""" cfg = _make_config(lora={"enable_lora": False}) - model, _ = model_creation_utils.from_pretrained(cfg, mesh=None, model_mode=model_creation_utils.MODEL_MODE_TRAIN) + model, _ = model_creation_utils.from_pretrained( + cfg, mesh=None, model_mode=model_creation_utils.MODEL_MODE_TRAIN) # Pydantic MaxTextConfig supports direct attribute access self.assertFalse(cfg.lora.enable_lora) result = lora_utils.apply_lora_to_model(model, None, cfg) @@ -192,7 +200,8 @@ def test_apply_lora_to_model_disabled(self): def test_apply_lora_to_model_adapters_loaded(self): """Test applying LoRA when adapters are already provided.""" cfg = _make_config(**{"lora_input_adapters_path": "some/path"}) - model, _ = model_creation_utils.from_pretrained(cfg, mesh=None, model_mode=model_creation_utils.MODEL_MODE_TRAIN) + model, _ = model_creation_utils.from_pretrained( + cfg, mesh=None, model_mode=model_creation_utils.MODEL_MODE_TRAIN) result = lora_utils.apply_lora_to_model(model, None, cfg) self.assertEqual(result, model) # is_lora_enabled checks for LoRAParam which Qwix adds. @@ -221,7 +230,8 @@ def _run_apply_lora_test( ) # Create a real small model using standard creation utils - model, mesh = model_creation_utils.from_pretrained(cfg, mesh=None, model_mode=model_creation_utils.MODEL_MODE_TRAIN) + model, mesh = model_creation_utils.from_pretrained( + cfg, mesh=None, model_mode=model_creation_utils.MODEL_MODE_TRAIN) # Verify model is NOT lora enabled initially self.assertFalse(lora_utils.is_lora_enabled(model)) @@ -245,20 +255,24 @@ def _run_apply_lora_test( self.assertTrue(lora_utils.is_lora_enabled(lora_model)) # Verify quantization if weight_qtype is set - flat_state_paths = [".".join(str(p) for p in path) for path, _ in state.flat_state()] + flat_state_paths = [ + ".".join(str(p) for p in path) for path, _ in state.flat_state() + ] if weight_qtype is not None: - self.assertTrue(any("qvalue" in path for path in flat_state_paths), "Expected quantized weights (qvalue) in state") + self.assertTrue(any("qvalue" in path for path in flat_state_paths), + "Expected quantized weights (qvalue) in state") else: - self.assertFalse( - any("qvalue" in path for path in flat_state_paths), "Did not expect quantized weights (qvalue) in state" - ) + self.assertFalse(any("qvalue" in path for path in flat_state_paths), + "Did not expect quantized weights (qvalue) in state") # Test fit for PeftTrainer trainer_cfg = peft_trainer.TrainingConfig(eval_every_n_steps=10) optimizer = optax.adam(1e-4) # This instantiation will fail if wrt=nnx.LoRAParam cannot find any matching params - trainer = peft_trainer.PeftTrainer(model=lora_model, optimizer=optimizer, training_config=trainer_cfg) + trainer = peft_trainer.PeftTrainer(model=lora_model, + optimizer=optimizer, + training_config=trainer_cfg) # Verify optimizer is indeed targeting LoRAParams opt_state = nnx.state(trainer.optimizer) @@ -274,11 +288,15 @@ def test_apply_lora_to_model_scan_layers_true(self): def test_apply_qlora_to_model_scan_layers_false(self): """Test applying QLoRA to model with scan_layers=False.""" - self._run_apply_lora_test(scan_layers=False, weight_qtype="int8", tile_size=32) + self._run_apply_lora_test(scan_layers=False, + weight_qtype="int8", + tile_size=32) def test_apply_qlora_to_model_scan_layers_true(self): """Test applying QLoRA to model with scan_layers=True.""" - self._run_apply_lora_test(scan_layers=True, weight_qtype="int8", tile_size=32) + self._run_apply_lora_test(scan_layers=True, + weight_qtype="int8", + tile_size=32) def test_apply_lora_multihost_mock(self): """Test applying LoRA with a dummy mesh to trigger the multi-host reshard callback.""" @@ -295,13 +313,17 @@ def test_restore_lora_from_path(self): }, scan_layers=False, ) - model, _ = model_creation_utils.from_pretrained(cfg, mesh=None, model_mode=model_creation_utils.MODEL_MODE_TRAIN) + model, _ = model_creation_utils.from_pretrained( + cfg, mesh=None, model_mode=model_creation_utils.MODEL_MODE_TRAIN) model = lora_utils.apply_lora_to_model(model, None, cfg) restored_state = nnx.state(model, nnx.LoRAParam) - with mock.patch("orbax.checkpoint.PyTreeCheckpointer.restore", return_value=restored_state) as mock_restore: - with mock.patch("flax.nnx.update") as mock_update: + with mock.patch("orbax.checkpoint.PyTreeCheckpointer.restore", + return_value=restored_state) as mock_restore: + with mock.patch( + "maxtext.utils.lora_utils.checkpointing._update_nnx_state_from_pure_dict" + ) as mock_update: lora_utils.restore_lora_from_path(model, cfg) mock_restore.assert_called_once() args, kwargs = mock_restore.call_args @@ -321,12 +343,17 @@ def test_sync_lora_metadata_default_syncs(self): "lora_restore_path": "dummy/path", "lora_rank": 0, "lora_alpha": 0.0, - } - ) + }) mock_metadata = mock.MagicMock() - mock_metadata.custom_metadata = {"lora": {"lora_rank": 32, "lora_alpha": 64.0}} + mock_metadata.custom_metadata = { + "lora": { + "lora_rank": 32, + "lora_alpha": 64.0 + } + } - with mock.patch("orbax.checkpoint.StandardCheckpointer.metadata", return_value=mock_metadata): + with mock.patch("orbax.checkpoint.StandardCheckpointer.metadata", + return_value=mock_metadata): lora_utils.sync_lora_metadata(cfg) self.assertEqual(cfg.lora.lora_rank, 32) self.assertEqual(cfg.lora.lora_alpha, 64.0) @@ -339,12 +366,17 @@ def test_sync_lora_metadata_matching_passes(self): "lora_restore_path": "dummy/path", "lora_rank": 32, "lora_alpha": 64.0, - } - ) + }) mock_metadata = mock.MagicMock() - mock_metadata.custom_metadata = {"lora": {"lora_rank": 32, "lora_alpha": 64.0}} + mock_metadata.custom_metadata = { + "lora": { + "lora_rank": 32, + "lora_alpha": 64.0 + } + } - with mock.patch("orbax.checkpoint.StandardCheckpointer.metadata", return_value=mock_metadata): + with mock.patch("orbax.checkpoint.StandardCheckpointer.metadata", + return_value=mock_metadata): # Should not raise ValueError lora_utils.sync_lora_metadata(cfg) self.assertEqual(cfg.lora.lora_rank, 32) @@ -358,13 +390,19 @@ def test_sync_lora_metadata_rank_mismatch_fails(self): "lora_restore_path": "dummy/path", "lora_rank": 8, "lora_alpha": 64.0, - } - ) + }) mock_metadata = mock.MagicMock() - mock_metadata.custom_metadata = {"lora": {"lora_rank": 32, "lora_alpha": 64.0}} + mock_metadata.custom_metadata = { + "lora": { + "lora_rank": 32, + "lora_alpha": 64.0 + } + } - with mock.patch("orbax.checkpoint.StandardCheckpointer.metadata", return_value=mock_metadata): - with self.assertRaisesRegex(ValueError, "Configured lora_rank .* does not match"): + with mock.patch("orbax.checkpoint.StandardCheckpointer.metadata", + return_value=mock_metadata): + with self.assertRaisesRegex(ValueError, + "Configured lora_rank .* does not match"): lora_utils.sync_lora_metadata(cfg) def test_sync_lora_metadata_alpha_mismatch_fails(self): @@ -375,26 +413,39 @@ def test_sync_lora_metadata_alpha_mismatch_fails(self): "lora_restore_path": "dummy/path", "lora_rank": 32, "lora_alpha": 16.0, - } - ) + }) mock_metadata = mock.MagicMock() - mock_metadata.custom_metadata = {"lora": {"lora_rank": 32, "lora_alpha": 64.0}} + mock_metadata.custom_metadata = { + "lora": { + "lora_rank": 32, + "lora_alpha": 64.0 + } + } - with mock.patch("orbax.checkpoint.StandardCheckpointer.metadata", return_value=mock_metadata): - with self.assertRaisesRegex(ValueError, "Configured lora_alpha .* does not match"): + with mock.patch("orbax.checkpoint.StandardCheckpointer.metadata", + return_value=mock_metadata): + with self.assertRaisesRegex(ValueError, + "Configured lora_alpha .* does not match"): lora_utils.sync_lora_metadata(cfg) def test_save_checkpoint_passes_metadata(self): """Test that save_checkpoint correctly generates and passes custom lora metadata to CheckpointManager.""" cfg = _make_config( - lora={"enable_lora": True, "lora_rank": 8, "lora_alpha": 16.0}, + lora={ + "enable_lora": True, + "lora_rank": 8, + "lora_alpha": 16.0 + }, enable_checkpointing=True, ) mock_manager = mock.MagicMock() mock_state = mock.MagicMock() with mock.patch("jax.block_until_ready"): - checkpointing.save_checkpoint(mock_manager, step=10, state=mock_state, config=cfg) + checkpointing.save_checkpoint(mock_manager, + step=10, + state=mock_state, + config=cfg) mock_manager.save.assert_called_once() _, kwargs = mock_manager.save.call_args self.assertIn("custom_metadata", kwargs) @@ -404,7 +455,11 @@ def test_save_and_restore_metadata_integration(self): """Integration test checking that Orbax CheckpointManager writes and reads custom LoRA metadata.""" cfg_save = _make_config( - lora={"enable_lora": True, "lora_rank": 8, "lora_alpha": 16.0}, + lora={ + "enable_lora": True, + "lora_rank": 8, + "lora_alpha": 16.0 + }, enable_checkpointing=True, ) @@ -420,7 +475,10 @@ def test_save_and_restore_metadata_integration(self): # Use save_checkpoint wrapper with a simple state dummy_state = {"weight": jnp.array([1.0, 2.0])} - checkpointing.save_checkpoint(manager, step=0, state=dummy_state, config=cfg_save) + checkpointing.save_checkpoint(manager, + step=0, + state=dummy_state, + config=cfg_save) manager.wait_until_finished() # Now verify that the saved checkpoint contains metadata on disk @@ -434,8 +492,7 @@ def test_save_and_restore_metadata_integration(self): "lora_restore_path": str(checkpoint_dir), "lora_rank": 0, "lora_alpha": 0.0, - } - ) + }) lora_utils.sync_lora_metadata(cfg_restore) # Verify values were successfully synced back @@ -483,7 +540,8 @@ def test_gemma4_lora_path_matching(self): ] for path in matching_paths: - self.assertTrue(compiled.search(path), f"Failed to match valid path: {path}") + self.assertTrue(compiled.search(path), + f"Failed to match valid path: {path}") # Expected non-matching paths (e.g. layernorm, embedding): non_matching_paths = [ @@ -496,7 +554,106 @@ def test_gemma4_lora_path_matching(self): ] for path in non_matching_paths: - self.assertFalse(compiled.search(path), f"Incorrectly matched invalid path: {path}") + self.assertFalse(compiled.search(path), + f"Incorrectly matched invalid path: {path}") + + def test_l2norm_pytree_void_quantized_arrays(self): + """Test that max_utils.l2norm_pytree safely skips non-numeric, void, or custom wrapper leaves.""" + from maxtext.utils import max_utils + + void_arr = np.frombuffer(b"\x00\x01\x02\x03", dtype="V4") + + class MockWithAux: + dtype = np.dtype("V4") + + tree = { + "float_param": jnp.array([3.0, 4.0], dtype=jnp.float32), + "int_param": jnp.array([1, 2], dtype=jnp.int8), + "quantized_void_param": void_arr, + "custom_with_aux_param": MockWithAux(), + } + norm = max_utils.l2norm_pytree(tree) + self.assertAlmostEqual(float(norm), 5.4772255, places=4) + + def test_other_lora_format_to_jax_format_mlp(self): + """Test that gate_proj, up_proj, and down_proj map to MaxText MLP modules.""" + raw_target_modules = [ + "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", + "down_proj" + ] + mapping = { + "q_proj": "query", + "k_proj": "key", + "v_proj": "value", + "o_proj": "out", + "gate_proj": "wi_0", + "up_proj": "wi_1", + "down_proj": "wo", + } + mapped = [mapping.get(m, m) for m in raw_target_modules] + self.assertIn("wi_0", mapped) + self.assertIn("wi_1", mapped) + self.assertIn("wo", mapped) + + def test_reshard_obj_to_mesh(self): + """Test reshard_obj_to_mesh with NNX state and variables.""" + + class DummyNNXModule(nnx.Module): + + def __init__(self): + self.w = nnx.Param(jnp.ones((4, 4))) + + mod = DummyNNXModule() + config = _make_config() + res = lora_utils.reshard_obj_to_mesh(mod, None, config) + self.assertIs(res, mod) + + def test_restore_qlora_base_weights(self): + qstate = nnx.State({ + "layer": { + "qvalue": jnp.array([1, 2], dtype=jnp.int8), + "scale": jnp.array([0.1, 0.2]), + } + }) + restored = lora_utils.restore_qlora_base_weights(qstate) + self.assertTrue(isinstance(restored["layer"], jax.ShapeDtypeStruct)) + + def test_restore_qlora_base_weights_edge_cases(self): + """Test restore_qlora_base_weights with missing keys, partial dicts, or non-dict structures.""" + # 1. Non-dict leaf (e.g. raw float array) should return as-is + raw_array = jnp.array([1.0, 2.0]) + np.testing.assert_array_equal( + lora_utils.restore_qlora_base_weights(raw_array), raw_array) + + # 2. Dict missing scale (only qvalue) should not convert to QArray + partial_qval = {"qvalue": jnp.array([1, 2], dtype=jnp.int8)} + res_qval = lora_utils.restore_qlora_base_weights(partial_qval) + self.assertNotIn("scale", res_qval) + + # 3. Dict missing qvalue (only scale) should not convert + partial_scale = {"scale": jnp.array([0.1, 0.2])} + res_scale = lora_utils.restore_qlora_base_weights(partial_scale) + self.assertNotIn("qvalue", res_scale) + + # 4. None input + self.assertIsNone(lora_utils.restore_qlora_base_weights(None)) + + def test_l2norm_pytree_non_numeric_and_bool_leaves(self): + """Test l2norm_pytree with boolean masks, string metadata, None, and empty PyTrees.""" + from maxtext.utils import max_utils + + tree = { + "float_val": jnp.array([3.0, 4.0], dtype=jnp.float32), + "bool_mask": jnp.array([True, False, True]), + "str_meta": "metadata_string", + "none_leaf": None, + "nested": { + "empty_list": [], + "another_float": jnp.array([0.0], dtype=jnp.float32), + }, + } + norm = max_utils.l2norm_pytree(tree) + self.assertAlmostEqual(float(norm), 5.0, places=4) if __name__ == "__main__": diff --git a/tests/unit/checkpointing_test.py b/tests/unit/checkpointing_test.py index f21f025175..e2360576f4 100644 --- a/tests/unit/checkpointing_test.py +++ b/tests/unit/checkpointing_test.py @@ -11,7 +11,6 @@ # 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. - """Unit tests for the checkpointing components.""" import asyncio @@ -47,7 +46,10 @@ def test_binary_chunked_stack(self): for shape in shapes: for num_tensors in [1, 2, 3, 5, 8, 12]: key = jax.random.PRNGKey(0) - tensors = [jax.random.normal(jax.random.fold_in(key, i), shape) for i in range(num_tensors)] + tensors = [ + jax.random.normal(jax.random.fold_in(key, i), shape) + for i in range(num_tensors) + ] # Test along various axes for axis in range(-len(shape) - 1, len(shape) + 1): @@ -62,8 +64,10 @@ class TensorHandlingTest(parameterized.TestCase): def setUp(self): super().setUp() self.mesh = Mesh(np.array(jax.devices()[:1]), axis_names=("x",)) - self.sharding_rank4 = NamedSharding(self.mesh, PartitionSpec("x", None, None, None)) - self.sharding_rank3 = NamedSharding(self.mesh, PartitionSpec("x", None, None)) + self.sharding_rank4 = NamedSharding(self.mesh, + PartitionSpec("x", None, None, None)) + self.sharding_rank3 = NamedSharding(self.mesh, + PartitionSpec("x", None, None)) def test_get_hf_loading_function_case_2_3_single_axis(self): # Tests Case 2/3 and lines 179 and gets loader for single axis stacked @@ -93,7 +97,8 @@ def getter_fn(key): hook_fn = None - loader_fn = get_hf_loading_function(hf_keys, getter_fn, hook_fn, target_leaf, config) + loader_fn = get_hf_loading_function(hf_keys, getter_fn, hook_fn, + target_leaf, config) result = loader_fn() @@ -134,7 +139,8 @@ def getter_fn(key): hook_fn = None - loader_fn = get_hf_loading_function(hf_keys, getter_fn, hook_fn, target_leaf, config) + loader_fn = get_hf_loading_function(hf_keys, getter_fn, hook_fn, + target_leaf, config) result = loader_fn() @@ -150,19 +156,22 @@ class LoadDynamicTest(parameterized.TestCase): @mock.patch("huggingface_hub.HfFileSystem") @mock.patch("google.cloud.storage.Client") - def test_build_gcs_cache_worker_cache_hit(self, mock_storage_client, mock_hf_fs): + def test_build_gcs_cache_worker_cache_hit(self, mock_storage_client, + mock_hf_fs): mock_client_instance = mock_storage_client.return_value mock_bucket = mock_client_instance.bucket.return_value mock_blob = mock_bucket.blob.return_value mock_blob.exists.return_value = True - load_dynamic.build_gcs_cache_worker("some_repo/model.safetensors", "gs://my-bucket/cache", "token") + load_dynamic.build_gcs_cache_worker("some_repo/model.safetensors", + "gs://my-bucket/cache", "token") mock_blob.exists.assert_called_once() mock_blob.upload_from_file.assert_not_called() @mock.patch("huggingface_hub.HfFileSystem") @mock.patch("google.cloud.storage.Client") - def test_build_gcs_cache_worker_cache_miss_success(self, mock_storage_client, mock_hf_fs): + def test_build_gcs_cache_worker_cache_miss_success(self, mock_storage_client, + mock_hf_fs): mock_fs_instance = mock_hf_fs.return_value mock_remote_file = mock.MagicMock() mock_fs_instance.open.return_value.__enter__.return_value = mock_remote_file @@ -172,13 +181,16 @@ def test_build_gcs_cache_worker_cache_miss_success(self, mock_storage_client, mo mock_blob = mock_bucket.blob.return_value mock_blob.exists.return_value = False - load_dynamic.build_gcs_cache_worker("some_repo/model.safetensors", "gs://my-bucket/cache", "token") + load_dynamic.build_gcs_cache_worker("some_repo/model.safetensors", + "gs://my-bucket/cache", "token") mock_blob.exists.assert_called_once() - mock_blob.upload_from_file.assert_called_once_with(mock_remote_file, client=mock_client_instance) + mock_blob.upload_from_file.assert_called_once_with( + mock_remote_file, client=mock_client_instance) @mock.patch("huggingface_hub.HfFileSystem") @mock.patch("google.cloud.storage.Client") - def test_build_gcs_cache_worker_retry_and_fail(self, mock_storage_client, mock_hf_fs): + def test_build_gcs_cache_worker_retry_and_fail(self, mock_storage_client, + mock_hf_fs): mock_fs_instance = mock_hf_fs.return_value mock_fs_instance.open.side_effect = Exception("Download failed") @@ -189,7 +201,8 @@ def test_build_gcs_cache_worker_retry_and_fail(self, mock_storage_client, mock_h with mock.patch("time.sleep"): with self.assertRaises(Exception): - load_dynamic.build_gcs_cache_worker("some_repo/model.safetensors", "gs://my-bucket/cache", "token") + load_dynamic.build_gcs_cache_worker("some_repo/model.safetensors", + "gs://my-bucket/cache", "token") @mock.patch.object(load_dynamic.huggingface_hub, "HfFileSystem") @mock.patch.object(load_dynamic.storage, "Client") @@ -236,7 +249,8 @@ def __init__(self): abstract_state = DummyAbstractState() path = "repo/meta-llama" - dummy_ret_val, loaded_vars = load_dynamic.load_safetensors_dynamic_state(path, abstract_state, config) + dummy_ret_val, loaded_vars = load_dynamic.load_safetensors_dynamic_state( + path, abstract_state, config) self.assertIsNone(dummy_ret_val) self.assertEqual(loaded_vars, {"params": {}}) @@ -262,7 +276,8 @@ def test_load_safetensors_dynamic_single_key(self): self.skipTest("SafetensorsLayout is not supported on Pathways backend.") # Save a single key (embedding weight) to a safetensors file dummy_weight = np.arange(1024, dtype=np.float32).reshape(256, 4) - safetensors.numpy.save_file({"model.embed_tokens.weight": dummy_weight}, str(self.safetensors_ckpt_path)) + safetensors.numpy.save_file({"model.embed_tokens.weight": dummy_weight}, + str(self.safetensors_ckpt_path)) # Setup mock config class MockConfig: @@ -280,13 +295,17 @@ def __init__(self): target_state = { "params": { "token_embedder": { - "embedding": jax.ShapeDtypeStruct(shape=(256, 4), dtype=np.float32, sharding=self.sharding) + "embedding": + jax.ShapeDtypeStruct(shape=(256, 4), + dtype=np.float32, + sharding=self.sharding) } } } abstract_state = train_state.TrainState.create( - apply_fn=lambda x: x, params=target_state["params"], tx=optax.identity() - ) + apply_fn=lambda x: x, + params=target_state["params"], + tx=optax.identity()) # Load using checkpointing framework dynamically loaded_data, loaded_vars = checkpointing.load_state_if_possible( @@ -316,15 +335,24 @@ class CheckpointMetadataTest(parameterized.TestCase): def test_load_checkpoint_metadata(self, mock_checkpointer_cls): mock_ckptr = mock_checkpointer_cls.return_value mock_metadata = mock.MagicMock() - mock_metadata.custom_metadata = {"lora": {"lora_rank": 8, "lora_alpha": 16.0}} + mock_metadata.custom_metadata = { + "lora": { + "lora_rank": 8, + "lora_alpha": 16.0 + } + } mock_ckptr.metadata.return_value = mock_metadata loaded_metadata = checkpointing.load_checkpoint_metadata("dummy/path") - self.assertEqual(loaded_metadata.get("lora"), {"lora_rank": 8, "lora_alpha": 16.0}) + self.assertEqual(loaded_metadata.get("lora"), { + "lora_rank": 8, + "lora_alpha": 16.0 + }) mock_ckptr.metadata.assert_called_once() @mock.patch.object(checkpointing.ocp, "StandardCheckpointer") - def test_load_checkpoint_metadata_handles_exceptions(self, mock_checkpointer_cls): + def test_load_checkpoint_metadata_handles_exceptions(self, + mock_checkpointer_cls): mock_ckptr = mock_checkpointer_cls.return_value mock_ckptr.metadata.side_effect = Exception("Checkpoint read error") @@ -341,6 +369,7 @@ def setUp(self): self.tmp_dir = epath.Path(self.create_tempdir().full_path) def test_save_restore_equivalence_single_item(self): + class FakeIterator: """A fake iterator for testing serialization.""" @@ -370,7 +399,8 @@ def __next__(self): handler.save(v0_path, item=iterator_v0) # v1 Save - wrapper = grain_utility.GrainCheckpointable(save_args=grain_utility.GrainCheckpointSave(item=iterator_v1)) + wrapper = grain_utility.GrainCheckpointable( + save_args=grain_utility.GrainCheckpointSave(item=iterator_v1)) class MockDirectory: """Mock directory for testing checkpointing.""" @@ -400,14 +430,15 @@ async def await_creation(self): # v1 Restore restored_iterator_v1 = FakeIterator(0) wrapper_restore = grain_utility.GrainCheckpointable( - restore_args=grain_utility.GrainCheckpointRestore(item=restored_iterator_v1) - ) + restore_args=grain_utility.GrainCheckpointRestore( + item=restored_iterator_v1)) load_func = asyncio.run(wrapper_restore.load(v1_path)) asyncio.run(load_func) self.assertEqual(restored_iterator_v1.state, 10) def test_save_restore_equivalence_list_item(self): + class FakeIterator: """A fake iterator for testing serialization.""" @@ -436,7 +467,8 @@ def set_state(self, state): handler.save(v0_path, item=item_v0) # v1 Save - wrapper = grain_utility.GrainCheckpointable(save_args=grain_utility.GrainCheckpointSave(item=item_v1)) + wrapper = grain_utility.GrainCheckpointable( + save_args=grain_utility.GrainCheckpointSave(item=item_v1)) class MockDirectory: """Mock directory for testing checkpointing.""" @@ -465,7 +497,9 @@ async def await_creation(self): # v0 Restore iterators_restore_v0 = [FakeIterator(0), FakeIterator(0)] - args_v0 = grain_utility.GrainCheckpointRestore(item=iterators_restore_v0, process_index=[0, 1], process_count=2) + args_v0 = grain_utility.GrainCheckpointRestore(item=iterators_restore_v0, + process_index=[0, 1], + process_count=2) handler.restore(v0_path, args=args_v0) self.assertEqual(iterators_restore_v0[0].state, 10) @@ -475,9 +509,7 @@ async def await_creation(self): iterators_restore_v1 = [FakeIterator(0), FakeIterator(0)] wrapper_restore = grain_utility.GrainCheckpointable( restore_args=grain_utility.GrainCheckpointRestore( - item=iterators_restore_v1, process_index=[0, 1], process_count=2 - ) - ) + item=iterators_restore_v1, process_index=[0, 1], process_count=2)) load_func = asyncio.run(wrapper_restore.load(v1_path)) asyncio.run(load_func) self.assertEqual(iterators_restore_v1[0].state, 10) diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py index ede6c1b35a..7ee094c038 100644 --- a/tests/unit/maxtext_utils_test.py +++ b/tests/unit/maxtext_utils_test.py @@ -1653,8 +1653,9 @@ def test_basic_abstraction(self): # Verify PartitionSpec was extracted correctly from the mock model's annotations # Path: params -> kernel -> spec + kernel_spec = annotations.params.kernel.get_value() if isinstance(annotations.params.kernel, nnx.Variable) else annotations.params.kernel self.assertEqual( - annotations.params.kernel.get_value(), + kernel_spec, PartitionSpec( "model", ), @@ -1669,7 +1670,7 @@ def test_shard_optimizer_over_data(self): # Original Pspec for optimizer was PartitionSpec(None). # add_data_to_sharding should find that dim 0 is compatible with mesh 'data' # and update it to PartitionSpec(('data',)). - opt_spec = annotations.optimizer.get_value() + opt_spec = annotations.optimizer.get_value() if isinstance(annotations.optimizer, nnx.Variable) else annotations.optimizer # Verify 'data' is now in the spec self.assertEqual(opt_spec, PartitionSpec(("data", "model"))) @@ -1681,11 +1682,11 @@ def test_optimizer_host_offload(self): _, _, shardings = maxtext_utils.get_abstract_state_nnx(self.config, self.mesh, self.nnx_init_trainstate_wrapper) # Optimizer state should be pinned to host - opt_sharding = shardings.optimizer.get_value() + opt_sharding = shardings.optimizer.get_value() if isinstance(shardings.optimizer, nnx.Variable) else shardings.optimizer self.assertEqual(opt_sharding.memory_kind, "pinned_host") # Params should still be on default memory (usually device) - param_sharding = shardings.params.kernel.get_value() + param_sharding = shardings.params.kernel.get_value() if isinstance(shardings.params.kernel, nnx.Variable) else shardings.params.kernel self.assertNotEqual(param_sharding.memory_kind, "pinned_host") def test_parameter_host_offload(self): @@ -1696,7 +1697,7 @@ def test_parameter_host_offload(self): _, _, shardings = maxtext_utils.get_abstract_state_nnx(self.config, self.mesh, self.nnx_init_trainstate_wrapper) # Parameters should be pinned to host - param_sharding = shardings.params.kernel.get_value() + param_sharding = shardings.params.kernel.get_value() if isinstance(shardings.params.kernel, nnx.Variable) else shardings.params.kernel self.assertEqual(param_sharding.memory_kind, "pinned_host") def test_invalid_init_fn(self): diff --git a/tests/unit/model_creation_utils_test.py b/tests/unit/model_creation_utils_test.py index b0ec54f694..f32a94e05b 100644 --- a/tests/unit/model_creation_utils_test.py +++ b/tests/unit/model_creation_utils_test.py @@ -11,7 +11,6 @@ # 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. - """Unit tests for model_creation_utils.py.""" import dataclasses @@ -122,7 +121,8 @@ def _run_fix(self, stored_shape, model_shape, sharding=None): arg = dataclasses.replace(arg, sharding=sharding) restore_args = {"kernel": arg} metadata_tree = {"kernel": _FakeArrayMetadata(shape=stored_shape)} - return _fix_restore_args_for_shape_mismatch(restore_args, metadata_tree, self.mesh) + return _fix_restore_args_for_shape_mismatch(restore_args, metadata_tree, + self.mesh) def test_scanned_ckpt_unscanned_model_raises_error(self): """Rank mismatch (scanned ckpt rank 4 vs unscanned model rank 3): arg must be unchanged.""" @@ -147,17 +147,19 @@ def test_same_rank_shape_mismatch_divisible_keeps_sharding(self): def test_same_rank_shape_mismatch_indivisible_falls_back_to_replicated(self): """Stored dim not divisible by mesh axis size: fall back to fully-replicated.""" if jax.device_count() < 2: - self.skipTest("Requires >=2 devices to construct an indivisible mesh axis.") + self.skipTest( + "Requires >=2 devices to construct an indivisible mesh axis.") devices = jax.local_devices()[:2] multi_mesh = jax.sharding.Mesh(devices, ("x",)) - sharding = jax.sharding.NamedSharding(multi_mesh, jax.sharding.PartitionSpec(None, "x", None)) + sharding = jax.sharding.NamedSharding( + multi_mesh, jax.sharding.PartitionSpec(None, "x", None)) # stored axis 1 = 3, mesh "x" = 2 -> 3 % 2 != 0 -> can't shard, must replicate. stored_shape = (256, 3, 128) model_shape = (256, 4, 128) arg = dataclasses.replace(_make_restore_arg(model_shape), sharding=sharding) fixed = _fix_restore_args_for_shape_mismatch( - {"kernel": arg}, {"kernel": _FakeArrayMetadata(shape=stored_shape)}, multi_mesh - ) + {"kernel": arg}, {"kernel": _FakeArrayMetadata(shape=stored_shape)}, + multi_mesh) arg = fixed["kernel"] # Replicated fallback: global_shape cleared, sharding swapped to fully-replicated. self.assertIsNone(arg.global_shape) @@ -200,13 +202,16 @@ def test_kv_heads_repeat_layout(self): # Distinguishable per-head values so the layout assertion is unambiguous. ckpt = jnp.array( [ - [[1.0, 1.0], [2.0, 2.0]], # embed=0: kv_head_0=[1,1], kv_head_1=[2,2] - [[3.0, 3.0], [4.0, 4.0]], # embed=1: kv_head_0=[3,3], kv_head_1=[4,4] + [[1.0, 1.0], [2.0, 2.0] + ], # embed=0: kv_head_0=[1,1], kv_head_1=[2,2] + [[3.0, 3.0], [4.0, 4.0] + ], # embed=1: kv_head_0=[3,3], kv_head_1=[4,4] ], dtype=jnp.float32, ) model = jnp.zeros((2, 4, 2), dtype=jnp.float32) - out = _align_checkpoint_to_model_shapes(ckpt, model, ("embed", "kv_heads", "kv_head_dim")) + out = _align_checkpoint_to_model_shapes( + ckpt, model, ("embed", "kv_heads", "kv_head_dim")) out_np = np.asarray(out) self.assertEqual(out_np.shape, (2, 4, 2)) # kv_heads axis layout for embed=0 must be [h0, h0, h1, h1]: @@ -219,31 +224,40 @@ def test_mlp_moe_zero_pads_at_end_wi_0(self): """mlp_moe axis must zero-pad at the end (not repeat) — wi_0/wi_1 last-axis case.""" ckpt = jnp.ones((2, 4, 6), dtype=jnp.float32) # exp=2, embed=4, mlp=6 model = jnp.zeros((2, 4, 8), dtype=jnp.float32) # mlp padded to 8 - out = _align_checkpoint_to_model_shapes(ckpt, model, ("exp", "embed_moe", "mlp_moe")) + out = _align_checkpoint_to_model_shapes(ckpt, model, + ("exp", "embed_moe", "mlp_moe")) out_np = np.asarray(out) self.assertEqual(out_np.shape, (2, 4, 8)) # First 6 cols preserved as ones; last 2 cols are zeros. - np.testing.assert_array_equal(out_np[..., :6], np.ones((2, 4, 6), dtype=np.float32)) - np.testing.assert_array_equal(out_np[..., 6:], np.zeros((2, 4, 2), dtype=np.float32)) + np.testing.assert_array_equal(out_np[..., :6], + np.ones((2, 4, 6), dtype=np.float32)) + np.testing.assert_array_equal(out_np[..., 6:], + np.zeros((2, 4, 2), dtype=np.float32)) def test_mlp_moe_zero_pads_at_end_wo(self): """wo has the MLP dim on axis 1 — zero-pad at the end of axis 1.""" ckpt = jnp.ones((2, 6, 4), dtype=jnp.float32) model = jnp.zeros((2, 8, 4), dtype=jnp.float32) - out = _align_checkpoint_to_model_shapes(ckpt, model, ("exp", "mlp_moe", "embed_moe")) + out = _align_checkpoint_to_model_shapes(ckpt, model, + ("exp", "mlp_moe", "embed_moe")) out_np = np.asarray(out) self.assertEqual(out_np.shape, (2, 8, 4)) - np.testing.assert_array_equal(out_np[:, :6, :], np.ones((2, 6, 4), dtype=np.float32)) - np.testing.assert_array_equal(out_np[:, 6:, :], np.zeros((2, 2, 4), dtype=np.float32)) + np.testing.assert_array_equal(out_np[:, :6, :], + np.ones((2, 6, 4), dtype=np.float32)) + np.testing.assert_array_equal(out_np[:, 6:, :], + np.zeros((2, 2, 4), dtype=np.float32)) def test_activation_mlp_zero_pads_bias(self): """activation_mlp axis (bias arrays) must also zero-pad at the end.""" ckpt = jnp.ones((2, 6), dtype=jnp.float32) model = jnp.zeros((2, 8), dtype=jnp.float32) - out = _align_checkpoint_to_model_shapes(ckpt, model, ("exp", "activation_mlp")) + out = _align_checkpoint_to_model_shapes(ckpt, model, + ("exp", "activation_mlp")) out_np = np.asarray(out) - np.testing.assert_array_equal(out_np[:, :6], np.ones((2, 6), dtype=np.float32)) - np.testing.assert_array_equal(out_np[:, 6:], np.zeros((2, 2), dtype=np.float32)) + np.testing.assert_array_equal(out_np[:, :6], + np.ones((2, 6), dtype=np.float32)) + np.testing.assert_array_equal(out_np[:, 6:], + np.zeros((2, 2), dtype=np.float32)) def test_unknown_axis_divisible_falls_back_to_repeat(self): """Unknown axis with divisible dims preserves prior repeat behavior (backward compat).""" @@ -256,7 +270,8 @@ def test_unknown_axis_non_divisible_raises(self): """Unknown axis with non-divisible dims must raise a clear error pointing at the axis.""" ckpt = jnp.ones((3,), dtype=jnp.float32) model = jnp.zeros((4,), dtype=jnp.float32) - with self.assertRaisesRegex(ValueError, "not registered in _VLLM_REPEAT_AXES"): + with self.assertRaisesRegex(ValueError, + "not registered in _VLLM_REPEAT_AXES"): _align_checkpoint_to_model_shapes(ckpt, model, ("unknown_axis",)) def test_logical_axes_none_falls_back_to_repeat(self): @@ -271,14 +286,16 @@ def test_kv_heads_non_divisible_raises(self): ckpt = jnp.ones((4, 3, 2), dtype=jnp.float32) model = jnp.zeros((4, 4, 2), dtype=jnp.float32) with self.assertRaisesRegex(ValueError, "kv_heads"): - _align_checkpoint_to_model_shapes(ckpt, model, ("embed", "kv_heads", "kv_head_dim")) + _align_checkpoint_to_model_shapes(ckpt, model, + ("embed", "kv_heads", "kv_head_dim")) def test_rank_mismatch_raises(self): """Different ranks (scanned vs unscanned ckpt) must raise the helpful message.""" ckpt = jnp.ones((4, 4, 2), dtype=jnp.float32) model = jnp.zeros((2, 4, 4, 2), dtype=jnp.float32) with self.assertRaisesRegex(ValueError, "different ranks"): - _align_checkpoint_to_model_shapes(ckpt, model, ("scan", "embed", "kv_heads", "kv_head_dim")) + _align_checkpoint_to_model_shapes( + ckpt, model, ("scan", "embed", "kv_heads", "kv_head_dim")) @pytest.mark.tpu_only @@ -290,8 +307,14 @@ class TestFuseMoeWeights(unittest.TestCase): def test_no_op_when_model_has_unfused_wi(self): """If model has wi_0/wi_1 (not fused), the helper returns the ckpt tree unchanged.""" - ckpt = {"wi_0": np.ones((2, 4, 3), dtype=np.float32), "wi_1": 2 * np.ones((2, 4, 3), dtype=np.float32)} - model = {"wi_0": np.zeros((2, 4, 3), dtype=np.float32), "wi_1": np.zeros((2, 4, 3), dtype=np.float32)} + ckpt = { + "wi_0": np.ones((2, 4, 3), dtype=np.float32), + "wi_1": 2 * np.ones((2, 4, 3), dtype=np.float32) + } + model = { + "wi_0": np.zeros((2, 4, 3), dtype=np.float32), + "wi_1": np.zeros((2, 4, 3), dtype=np.float32) + } out = _fuse_moe_weights(ckpt, model) self.assertIn("wi_0", out) self.assertIn("wi_1", out) @@ -324,18 +347,39 @@ def test_fuse_with_mlp_padding_pads_each_half(self): self.assertEqual(out["wi"].shape, (2, 4, 8)) # First half: wi_0_data (3 cols) then 1 col of zeros. np.testing.assert_array_equal(out["wi"][..., :3], wi_0) - np.testing.assert_array_equal(out["wi"][..., 3:4], np.zeros((2, 4, 1), dtype=np.float32)) + np.testing.assert_array_equal(out["wi"][..., 3:4], + np.zeros((2, 4, 1), dtype=np.float32)) # Second half: wi_1_data (3 cols) then 1 col of zeros. np.testing.assert_array_equal(out["wi"][..., 4:7], wi_1) - np.testing.assert_array_equal(out["wi"][..., 7:8], np.zeros((2, 4, 1), dtype=np.float32)) + np.testing.assert_array_equal(out["wi"][..., 7:8], + np.zeros((2, 4, 1), dtype=np.float32)) def test_recurses_through_nested_dicts(self): """The helper should walk arbitrary nested dicts and only fuse at MoE blocks.""" wi_0 = np.ones((2, 4, 3), dtype=np.float32) wi_1 = 2 * np.ones((2, 4, 3), dtype=np.float32) other = np.full((2, 4), 7.0, dtype=np.float32) - ckpt = {"layers": {"0": {"moe": {"wi_0": wi_0, "wi_1": wi_1}, "other": other}}} - model = {"layers": {"0": {"moe": {"wi": np.zeros((2, 4, 6), dtype=np.float32)}, "other": other}}} + ckpt = { + "layers": { + "0": { + "moe": { + "wi_0": wi_0, + "wi_1": wi_1 + }, + "other": other + } + } + } + model = { + "layers": { + "0": { + "moe": { + "wi": np.zeros((2, 4, 6), dtype=np.float32) + }, + "other": other + } + } + } out = _fuse_moe_weights(ckpt, model) self.assertEqual(out["layers"]["0"]["moe"]["wi"].shape, (2, 4, 6)) np.testing.assert_array_equal(out["layers"]["0"]["other"], other) @@ -356,22 +400,27 @@ class TestShardedZeroPad(unittest.TestCase): def test_evenly_shardable_true_when_stored_dim_divides_mesh_axis(self): devices = jax.local_devices()[:1] mesh = jax.sharding.Mesh(devices, ("model",)) - sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "model", None)) + sharding = jax.sharding.NamedSharding( + mesh, jax.sharding.PartitionSpec(None, "model", None)) restore_arg = _make_restore_arg((128, 1024, 2048)) restore_arg = dataclasses.replace(restore_arg, sharding=sharding) # 1-device mesh: any stored dim trivially divides 1. - self.assertTrue(_stored_shape_evenly_shardable(restore_arg, (128, 768, 2048))) + self.assertTrue( + _stored_shape_evenly_shardable(restore_arg, (128, 768, 2048))) def test_evenly_shardable_false_when_stored_dim_indivisible(self): if jax.device_count() < 2: - self.skipTest("Requires >=2 devices to construct a non-trivial mesh size.") + self.skipTest( + "Requires >=2 devices to construct a non-trivial mesh size.") devices = jax.local_devices()[:2] mesh = jax.sharding.Mesh(devices, ("model",)) - sharding = jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec(None, "model", None)) + sharding = jax.sharding.NamedSharding( + mesh, jax.sharding.PartitionSpec(None, "model", None)) restore_arg = _make_restore_arg((128, 8, 2048)) restore_arg = dataclasses.replace(restore_arg, sharding=sharding) # stored_shape[1]=3 is not divisible by mesh "model"=2 -> must fall back to replicated. - self.assertFalse(_stored_shape_evenly_shardable(restore_arg, (128, 3, 2048))) + self.assertFalse(_stored_shape_evenly_shardable(restore_arg, + (128, 3, 2048))) def test_zero_pad_axis_no_op_when_extra_zero(self): arr = jnp.ones((2, 4, 2), dtype=jnp.float32) @@ -383,8 +432,10 @@ def test_zero_pad_axis_unsharded_falls_back_to_global_pad(self): out = _zero_pad_axis(arr, axis=1, extra=2) out_np = np.asarray(out) self.assertEqual(out_np.shape, (2, 6, 2)) - np.testing.assert_array_equal(out_np[:, :4, :], np.ones((2, 4, 2), dtype=np.float32)) - np.testing.assert_array_equal(out_np[:, 4:, :], np.zeros((2, 2, 2), dtype=np.float32)) + np.testing.assert_array_equal(out_np[:, :4, :], + np.ones((2, 4, 2), dtype=np.float32)) + np.testing.assert_array_equal(out_np[:, 4:, :], + np.zeros((2, 2, 2), dtype=np.float32)) def test_zero_pad_axis_sharded_pads_each_local_shard(self): """Per-shard layout: zeros land at the tail of *each* shard, not the global tail. @@ -395,7 +446,8 @@ def test_zero_pad_axis_sharded_pads_each_local_shard(self): tails are zero, per-shard heads preserve the original ckpt slice. """ if jax.device_count() < 4: - self.skipTest("Requires >=4 devices to exercise the sharded pad path with TP=4.") + self.skipTest( + "Requires >=4 devices to exercise the sharded pad path with TP=4.") devices = jax.local_devices()[:4] mesh = jax.sharding.Mesh(devices, ("model",)) spec = jax.sharding.PartitionSpec(None, "model", None) @@ -404,16 +456,19 @@ def test_zero_pad_axis_sharded_pads_each_local_shard(self): ckpt_np = np.arange(2 * 8 * 2, dtype=np.float32).reshape(2, 8, 2) arr = jax.device_put(jnp.asarray(ckpt_np), sharding) - out = _zero_pad_axis(arr, axis=1, extra=4) # 8 -> 12, 4 shards -> +1 zero per shard + out = _zero_pad_axis(arr, axis=1, + extra=4) # 8 -> 12, 4 shards -> +1 zero per shard out_np = np.asarray(out) self.assertEqual(out_np.shape, (2, 12, 2)) for s in range(4): # Per-shard head holds the original ckpt slice [s*2:s*2+2]. for i in range(2): - np.testing.assert_array_equal(out_np[:, s * 3 + i, :], ckpt_np[:, s * 2 + i, :]) + np.testing.assert_array_equal(out_np[:, s * 3 + i, :], + ckpt_np[:, s * 2 + i, :]) # Per-shard tail is zero. - np.testing.assert_array_equal(out_np[:, s * 3 + 2, :], np.zeros((2, 2), dtype=np.float32)) + np.testing.assert_array_equal(out_np[:, s * 3 + 2, :], + np.zeros((2, 2), dtype=np.float32)) @pytest.mark.tpu_only @@ -430,6 +485,7 @@ class TestPartitionSpecUnwrapForAlignment(unittest.TestCase): """ def test_unwrap_yields_partition_spec_and_enables_zero_pad(self): + class _TinyMoE(nnx.Module): def __init__(self, rngs: nnx.Rngs): @@ -441,7 +497,8 @@ def __init__(self, rngs: nnx.Rngs): # flax >= 0.12 requires an active mesh + logical_axis_rules to construct # a Variable annotated with logical sharding names. devices = jax.local_devices()[:1] - mesh = jax.sharding.Mesh(devices, ("x",), axis_types=(jax.sharding.AxisType.Explicit,)) + mesh = jax.sharding.Mesh(devices, ("x",), + axis_types=(jax.sharding.AxisType.Explicit,)) rules = (("exp", None), ("mlp_moe", None), ("embed_moe", None)) with jax.sharding.set_mesh(mesh), nn.logical_axis_rules(rules): module = nnx.eval_shape(lambda: _TinyMoE(nnx.Rngs(params=0))) @@ -468,8 +525,10 @@ def __init__(self, rngs: nnx.Rngs): out = _align_checkpoint_to_model_shapes(ckpt, model, wo_axes) out_np = np.asarray(out) self.assertEqual(out_np.shape, (2, 6, 4)) - np.testing.assert_array_equal(out_np[:, :4, :], np.ones((2, 4, 4), dtype=np.float32)) - np.testing.assert_array_equal(out_np[:, 4:, :], np.zeros((2, 2, 4), dtype=np.float32)) + np.testing.assert_array_equal(out_np[:, :4, :], + np.ones((2, 4, 4), dtype=np.float32)) + np.testing.assert_array_equal(out_np[:, 4:, :], + np.zeros((2, 2, 4), dtype=np.float32)) class TestGetTransformerModel(unittest.TestCase): @@ -481,26 +540,35 @@ def setUp(self): def test_returns_linen_module_when_rngs_is_none(self): """Without rngs, should return a Linen nn.Module.""" - model = model_creation_utils.get_transformer_model(self.config, self.mesh, quant=None, rngs=None) + model = model_creation_utils.get_transformer_model(self.config, + self.mesh, + quant=None, + rngs=None) self.assertIsInstance(model, nn.Module) def test_returns_nnx_module_when_rngs_provided(self): """With rngs, should return an NNX nnx.Module.""" - model = nnx.eval_shape( - lambda: model_creation_utils.get_transformer_model( - self.config, self.mesh, quant=None, rngs=nnx.Rngs(params=0, dropout=1, aqt=2) - ) - ) + model = nnx.eval_shape(lambda: model_creation_utils.get_transformer_model( + self.config, + self.mesh, + quant=None, + rngs=nnx.Rngs(params=0, dropout=1, aqt=2))) self.assertIsInstance(model, nnx.Module) def test_respects_model_mode_prefill(self): """Linen model created with MODEL_MODE_PREFILL should differ from train mode.""" linen_train = model_creation_utils.get_transformer_model( - self.config, self.mesh, quant=None, model_mode=MODEL_MODE_TRAIN, rngs=None - ) + self.config, + self.mesh, + quant=None, + model_mode=MODEL_MODE_TRAIN, + rngs=None) linen_prefill = model_creation_utils.get_transformer_model( - self.config, self.mesh, quant=None, model_mode=MODEL_MODE_PREFILL, rngs=None - ) + self.config, + self.mesh, + quant=None, + model_mode=MODEL_MODE_PREFILL, + rngs=None) # Both are still nn.Module instances self.assertIsInstance(linen_train, nn.Module) self.assertIsInstance(linen_prefill, nn.Module) @@ -518,9 +586,8 @@ def test_returns_linen_model_without_rngs(self): self.assertIsInstance(model, nn.Module) def test_returns_nnx_model_with_rngs(self): - model = nnx.eval_shape( - lambda: model_creation_utils.create_model(self.config, self.mesh, rngs=nnx.Rngs(params=0, dropout=1, aqt=2)) - ) + model = nnx.eval_shape(lambda: model_creation_utils.create_model( + self.config, self.mesh, rngs=nnx.Rngs(params=0, dropout=1, aqt=2))) self.assertIsInstance(model, nnx.Module) def test_model_mode_train_default(self): @@ -538,24 +605,31 @@ def setUp(self): def test_linen_path_rngs_none(self): """from_config with rngs=None should return a Linen nn.Module.""" - model = model_creation_utils.from_config(self.config, mesh=self.mesh, rngs=None) + model = model_creation_utils.from_config(self.config, + mesh=self.mesh, + rngs=None) self.assertIsInstance(model, nn.Module) def test_nnx_path_with_rngs(self): """from_config with rngs provided should return an NNX nnx.Module.""" - model = nnx.eval_shape( - lambda: model_creation_utils.from_config(self.config, mesh=self.mesh, rngs=nnx.Rngs(params=0, dropout=1, aqt=2)) - ) + model = nnx.eval_shape(lambda: model_creation_utils.from_config( + self.config, mesh=self.mesh, rngs=nnx.Rngs(params=0, dropout=1, aqt=2))) self.assertIsInstance(model, nnx.Module) def test_mesh_created_from_devices_when_none(self): """from_config should work when mesh is None (creates mesh internally).""" - model = model_creation_utils.from_config(self.config, devices=None, mesh=None, rngs=None) + model = model_creation_utils.from_config(self.config, + devices=None, + mesh=None, + rngs=None) self.assertIsInstance(model, nn.Module) def test_model_mode_is_forwarded(self): """from_config should accept and forward model_mode.""" - model = model_creation_utils.from_config(self.config, mesh=self.mesh, model_mode=MODEL_MODE_PREFILL, rngs=None) + model = model_creation_utils.from_config(self.config, + mesh=self.mesh, + model_mode=MODEL_MODE_PREFILL, + rngs=None) self.assertIsInstance(model, nn.Module) def test_explicit_shard_mode_creates_mesh_with_explicit_axis_types(self): @@ -574,13 +648,15 @@ def setUp(self): self.mesh = _make_mesh(self.config) def test_returns_tuple_of_callable_and_module(self): - create_fn, abstract_model = model_creation_utils.create_nnx_abstract_model(self.config, mesh=self.mesh) + create_fn, abstract_model = model_creation_utils.create_nnx_abstract_model( + self.config, mesh=self.mesh) self.assertTrue(callable(create_fn)) self.assertIsInstance(abstract_model, nnx.Module) def test_abstract_model_has_abstract_arrays(self): """Abstract model leaves should be ShapeDtypeStruct, not concrete arrays.""" - _, abstract_model = model_creation_utils.create_nnx_abstract_model(self.config, mesh=self.mesh) + _, abstract_model = model_creation_utils.create_nnx_abstract_model( + self.config, mesh=self.mesh) _, state = nnx.split(abstract_model) leaves = jax.tree.leaves(state) self.assertGreater(len(leaves), 0) @@ -591,7 +667,8 @@ def test_abstract_model_has_abstract_arrays(self): def test_create_fn_produces_concrete_model(self): """The returned create_fn should produce a real (concrete) NNX Module.""" - create_fn, _ = model_creation_utils.create_nnx_abstract_model(self.config, mesh=self.mesh) + create_fn, _ = model_creation_utils.create_nnx_abstract_model( + self.config, mesh=self.mesh) with self.mesh: concrete = create_fn() self.assertIsInstance(concrete, nnx.Module) @@ -601,7 +678,8 @@ def test_create_fn_produces_concrete_model(self): def test_works_without_explicit_mesh(self): """create_nnx_abstract_model should work when mesh=None (from_config creates mesh).""" - create_fn, abstract_model = model_creation_utils.create_nnx_abstract_model(self.config, mesh=None) + create_fn, abstract_model = model_creation_utils.create_nnx_abstract_model( + self.config, mesh=None) self.assertTrue(callable(create_fn)) self.assertIsInstance(abstract_model, nnx.Module) @@ -609,16 +687,14 @@ def test_explicit_rng_key_is_used(self): """Passing a rng_key should not raise and returns valid abstract model.""" rng_key = jax.random.PRNGKey(42) create_fn, abstract_model = model_creation_utils.create_nnx_abstract_model( - self.config, mesh=self.mesh, rng_key=rng_key - ) + self.config, mesh=self.mesh, rng_key=rng_key) self.assertTrue(callable(create_fn)) self.assertIsInstance(abstract_model, nnx.Module) def test_prefill_model_mode(self): """create_nnx_abstract_model should accept MODEL_MODE_PREFILL.""" _, abstract_model = model_creation_utils.create_nnx_abstract_model( - self.config, mesh=self.mesh, model_mode=MODEL_MODE_PREFILL - ) + self.config, mesh=self.mesh, model_mode=MODEL_MODE_PREFILL) self.assertIsInstance(abstract_model, nnx.Module) @@ -643,12 +719,16 @@ def test_mesh_none_uses_abstract_model_mesh(self): def test_explicit_rng_key(self): """An explicit rng_key should be accepted without error.""" rng_key = jax.random.PRNGKey(99) - model = model_creation_utils.from_pretrained(self.config, self.mesh, rng_key=rng_key) + model = model_creation_utils.from_pretrained(self.config, + self.mesh, + rng_key=rng_key) self.assertIsInstance(model, models.Transformer) def test_inference_mode_disables_dropout_rng(self): """MODEL_MODE_PREFILL should create rngs without a dropout key.""" - model = model_creation_utils.from_pretrained(self.config, self.mesh, model_mode=MODEL_MODE_PREFILL) + model = model_creation_utils.from_pretrained(self.config, + self.mesh, + model_mode=MODEL_MODE_PREFILL) self.assertIsInstance(model, models.Transformer) def test_debug_sharding_flag(self): @@ -688,7 +768,8 @@ def test_load_nnx_checkpoint(self, mock_ocp): mock_ocp.checkpoint_utils.construct_restore_args.return_value = {} mock_ocp.ArrayRestoreArgs = ocp.ArrayRestoreArgs - cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/nnx_ckpt") + cfg = _make_config(enable_checkpointing=True, + load_parameters_path="gs://fake/nnx_ckpt") model = model_creation_utils.from_pretrained(cfg, self.mesh) self.assertIsInstance(model, models.Transformer) @@ -707,7 +788,8 @@ def test_load_linen_checkpoint(self, mock_ocp): mock_ocp.checkpoint_utils.construct_restore_args.return_value = {} mock_ocp.ArrayRestoreArgs = ocp.ArrayRestoreArgs - cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/linen_ckpt") + cfg = _make_config(enable_checkpointing=True, + load_parameters_path="gs://fake/linen_ckpt") model = model_creation_utils.from_pretrained(cfg, self.mesh) self.assertIsInstance(model, models.Transformer) @@ -719,18 +801,21 @@ def test_checkpoint_load_error_propagates(self, mock_ocp): mock_ocp.Checkpointer.return_value = mock_ckptr mock_ocp.PyTreeCheckpointHandler.return_value = MagicMock() - cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/bad_ckpt") + cfg = _make_config(enable_checkpointing=True, + load_parameters_path="gs://fake/bad_ckpt") with self.assertRaises(RuntimeError): model_creation_utils.from_pretrained(cfg, self.mesh) - @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + @patch( + "maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata" + ) def test_scan_layers_mismatch_raises_error(self, mock_load_meta): """ValueError is raised if run specifies scan_layers=True but checkpoint specifies scan_layers=False.""" mock_load_meta.return_value = {"scan_layers": False} - cfg = _make_config( - enable_checkpointing=True, load_parameters_path="gs://fake/scan_layers_false_ckpt", scan_layers=True - ) + cfg = _make_config(enable_checkpointing=True, + load_parameters_path="gs://fake/scan_layers_false_ckpt", + scan_layers=True) with self.assertRaises(ValueError) as context: model_creation_utils.from_pretrained(cfg, self.mesh) @@ -740,7 +825,9 @@ def test_scan_layers_mismatch_raises_error(self, mock_load_meta): str(context.exception), ) - @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + @patch( + "maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata" + ) @patch("maxtext.utils.model_creation_utils.ocp") def test_scan_layers_match_no_error(self, mock_ocp, mock_load_meta): """If the run specifies scan_layers=True and the checkpoint matches, it proceeds without error.""" @@ -754,16 +841,19 @@ def test_scan_layers_match_no_error(self, mock_ocp, mock_load_meta): mock_ocp.checkpoint_utils.construct_restore_args.return_value = {} mock_ocp.ArrayRestoreArgs = ocp.ArrayRestoreArgs - cfg = _make_config( - enable_checkpointing=True, load_parameters_path="gs://fake/scan_layers_true_ckpt", scan_layers=True - ) + cfg = _make_config(enable_checkpointing=True, + load_parameters_path="gs://fake/scan_layers_true_ckpt", + scan_layers=True) model = model_creation_utils.from_pretrained(cfg, self.mesh) self.assertIsInstance(model, models.Transformer) - @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + @patch( + "maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata" + ) @patch("maxtext.utils.model_creation_utils.ocp") - def test_scan_layers_missing_metadata_no_error(self, mock_ocp, mock_load_meta): + def test_scan_layers_missing_metadata_no_error(self, mock_ocp, + mock_load_meta): """Skip verification and proceed if custom_metadata lacks 'scan_layers'.""" mock_load_meta.return_value = {} @@ -776,8 +866,9 @@ def test_scan_layers_missing_metadata_no_error(self, mock_ocp, mock_load_meta): mock_ocp.ArrayRestoreArgs = ocp.ArrayRestoreArgs cfg = _make_config( - enable_checkpointing=True, load_parameters_path="gs://fake/scan_layers_missing_ckpt", scan_layers=True - ) + enable_checkpointing=True, + load_parameters_path="gs://fake/scan_layers_missing_ckpt", + scan_layers=True) model = model_creation_utils.from_pretrained(cfg, self.mesh) self.assertIsInstance(model, models.Transformer) @@ -795,15 +886,18 @@ def test_returns_linen_train_state_and_annotations(self): """Should return a linen TrainState whose params mirror the NNX model's nnx.Param values.""" # Build a real (small) NNX model WITHOUT any patch active so from_pretrained # runs normally and produces concrete jax.Array weights. - real_nnx_model = model_creation_utils.from_pretrained(self.config, mesh=self.mesh) + real_nnx_model = model_creation_utils.from_pretrained(self.config, + mesh=self.mesh) - linen_model = model_creation_utils.from_config(self.config, mesh=self.mesh, rngs=None) + linen_model = model_creation_utils.from_config(self.config, + mesh=self.mesh, + rngs=None) # Now patch from_pretrained so setup_decode_state_from_nnx never touches a checkpoint. - with patch("maxtext.utils.model_creation_utils.from_pretrained", return_value=real_nnx_model) as mock_fp: + with patch("maxtext.utils.model_creation_utils.from_pretrained", + return_value=real_nnx_model) as mock_fp: state, state_mesh_annotations = model_creation_utils.setup_decode_state_from_nnx( - linen_model, self.config, self.rng, self.mesh - ) + linen_model, self.config, self.rng, self.mesh) # from_pretrained must have been called with the right model_mode. mock_fp.assert_called_once() @@ -838,12 +932,15 @@ class TestVerifyAndSyncScanLayers(unittest.TestCase): def setUp(self): self.mesh = Mesh(np.array(jax.devices()[:1]), axis_names=("x",)) - @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + @patch( + "maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata" + ) def test_sync_to_false_when_implicit(self, mock_load_meta): """If scan_layers is not explicit, sync scan_layers to False from checkpoint metadata.""" mock_load_meta.return_value = {"scan_layers": False} # Create config without scan_layers in kwargs so it's not explicit - cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt") + cfg = _make_config(enable_checkpointing=True, + load_parameters_path="gs://fake/ckpt") # Pre-assertions pydantic_cfg = getattr(cfg, "_pydantic_config", cfg) @@ -858,20 +955,27 @@ def test_sync_to_false_when_implicit(self, mock_load_meta): synced_pydantic_cfg = getattr(synced_cfg, "_pydantic_config", synced_cfg) self.assertFalse(synced_pydantic_cfg.scan_layers) - @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + @patch( + "maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata" + ) def test_sync_to_true_when_implicit(self, mock_load_meta): """If scan_layers is not explicit, sync scan_layers to True from checkpoint metadata.""" mock_load_meta.return_value = {"scan_layers": True} - cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt") + cfg = _make_config(enable_checkpointing=True, + load_parameters_path="gs://fake/ckpt") synced_cfg = model_creation_utils.verify_and_sync_scan_layers(cfg) self.assertTrue(synced_cfg.scan_layers) - @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + @patch( + "maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata" + ) def test_explicit_match_raises_no_error(self, mock_load_meta): """If scan_layers is explicit and matches checkpoint metadata, no error is raised.""" mock_load_meta.return_value = {"scan_layers": True} - cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt", scan_layers=True) + cfg = _make_config(enable_checkpointing=True, + load_parameters_path="gs://fake/ckpt", + scan_layers=True) # Pre-assertions pydantic_cfg = getattr(cfg, "_pydantic_config", cfg) @@ -881,11 +985,15 @@ def test_explicit_match_raises_no_error(self, mock_load_meta): synced_cfg = model_creation_utils.verify_and_sync_scan_layers(cfg) self.assertTrue(synced_cfg.scan_layers) - @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + @patch( + "maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata" + ) def test_explicit_mismatch_raises_value_error(self, mock_load_meta): """If scan_layers is explicit and mismatches checkpoint metadata, ValueError is raised.""" mock_load_meta.return_value = {"scan_layers": False} - cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt", scan_layers=True) + cfg = _make_config(enable_checkpointing=True, + load_parameters_path="gs://fake/ckpt", + scan_layers=True) # Pre-assertions pydantic_cfg = getattr(cfg, "_pydantic_config", cfg) @@ -897,14 +1005,17 @@ def test_explicit_mismatch_raises_value_error(self, mock_load_meta): self.assertIn("Configuration mismatch", str(context.exception)) - @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + @patch( + "maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata" + ) @patch("maxtext.utils.model_creation_utils.max_logging.log") def test_sync_log_and_early_return(self, mock_log, mock_load_meta): """Test that we log only when auto-resolution actually changes scan_layers, and early return otherwise.""" # Scenario A: saved_scan_layers == config.scan_layers (default True). # Since they match, it should return early, and NO log should be printed. mock_load_meta.return_value = {"scan_layers": True} - cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt") + cfg = _make_config(enable_checkpointing=True, + load_parameters_path="gs://fake/ckpt") synced_cfg = model_creation_utils.verify_and_sync_scan_layers(cfg) self.assertTrue(synced_cfg.scan_layers) @@ -913,11 +1024,13 @@ def test_sync_log_and_early_return(self, mock_log, mock_load_meta): # Scenario B: saved_scan_layers (False) != config.scan_layers (default True), and is not explicit (implicit). # It should log the auto-resolution and update scan_layers to False. mock_load_meta.return_value = {"scan_layers": False} - cfg2 = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt") + cfg2 = _make_config(enable_checkpointing=True, + load_parameters_path="gs://fake/ckpt") synced_cfg2 = model_creation_utils.verify_and_sync_scan_layers(cfg2) self.assertFalse(synced_cfg2.scan_layers) - mock_log.assert_called_once_with("Setting scan_layers=False loaded from checkpoint metadata.") + mock_log.assert_called_once_with( + "Setting scan_layers=False loaded from checkpoint metadata.") if __name__ == "__main__": @@ -937,7 +1050,8 @@ def _make_nnx_metadata_mock(self): @patch("maxtext.utils.model_creation_utils.subprocess.run") @patch("maxtext.utils.model_creation_utils.get_token") @patch("maxtext.utils.model_creation_utils.epath.Path") - def test_auth_success_with_config_token(self, mock_path, mock_get_token, mock_run, mock_ocp): + def test_auth_success_with_config_token(self, mock_path, mock_get_token, + mock_run, mock_ocp): config = _make_config( convert_checkpoint_if_possible=True, base_output_directory="gs://fake_bucket/fake_run", @@ -968,7 +1082,8 @@ def test_auth_success_with_config_token(self, mock_path, mock_get_token, mock_ru @patch("maxtext.utils.model_creation_utils.subprocess.run") @patch("maxtext.utils.model_creation_utils.get_token") @patch("maxtext.utils.model_creation_utils.epath.Path") - def test_auth_success_with_cached_token(self, mock_path, mock_get_token, mock_run, mock_ocp): + def test_auth_success_with_cached_token(self, mock_path, mock_get_token, + mock_run, mock_ocp): config = _make_config( convert_checkpoint_if_possible=True, base_output_directory="gs://fake_bucket/fake_run", @@ -1001,7 +1116,8 @@ def test_auth_success_with_cached_token(self, mock_path, mock_get_token, mock_ru @patch("maxtext.utils.model_creation_utils.subprocess.run") @patch("maxtext.utils.model_creation_utils.get_token") @patch("maxtext.utils.model_creation_utils.epath.Path") - def test_auth_failure_no_token(self, mock_path, mock_get_token, mock_run, mock_ocp): + def test_auth_failure_no_token(self, mock_path, mock_get_token, mock_run, + mock_ocp): config = _make_config( convert_checkpoint_if_possible=True, base_output_directory="gs://fake_bucket/fake_run",