Skip to content

Commit 558f0f2

Browse files
RexBearIURexBearIU
committed
feat(nnx): support native Flax NNX PEFT/LoRA training loop with pre-trained base weight restoration and optimized upfront conversion
- Implemented standard public `v.set_value()` API for lora checkpoint restoration and complete validation fixes. - Integrated dynamic quantization, sharding propagation, and merge overlays into native checkpointing pipeline. - Added comprehensive E2E validation scripts for pre_train, sft_native, and sft_custom paths verifying loss continuity and standalone restore. - Standardized parameter sharding specifications and resharding helpers for safe multi-node TPU execution and Zero-1/FSDP alignment.
1 parent d557e88 commit 558f0f2

20 files changed

Lines changed: 2953 additions & 396 deletions

‎.github/workflows/nnx_moe_qlora_e2e.yml‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# https://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
name: Flax NNX QLoRA MoE E2E Tests
16+
17+
on:
18+
pull_request:
19+
branches:
20+
- nnx-lora-support
21+
- main
22+
paths:
23+
- 'src/maxtext/utils/lora_utils.py'
24+
- 'src/maxtext/trainers/**'
25+
- 'src/maxtext/configs/**'
26+
workflow_dispatch:
27+
28+
permissions:
29+
contents: read
30+
31+
jobs:
32+
run-moe-qlora-e2e:
33+
name: Run MoE QLoRA E2E Tests
34+
runs-on: self-hosted-tpu-v6e
35+
steps:
36+
- name: Checkout Repository
37+
uses: actions/checkout@v4
38+
39+
- name: Setup Environment
40+
run: |
41+
# Use local venv or install dependencies
42+
if [ -d ".venv" ]; then
43+
source .venv/bin/activate
44+
else
45+
python3 -m venv .venv
46+
source .venv/bin/activate
47+
pip install --upgrade pip
48+
pip install -e .
49+
fi
50+
51+
- name: Run E2E Test Runner Script
52+
run: |
53+
./run_e2e_moe_qlora_tests.sh

‎run_e2e_llama31_qlora_tests.sh‎

Lines changed: 242 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,242 @@
1+
#!/usr/bin/env bash
2+
# Copyright 2026 Google LLC
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# https://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
16+
set -o pipefail
17+
18+
# Detect directories
19+
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
20+
WORKSPACE_DIR="${SCRIPT_DIR}"
21+
22+
VENV_PYTHON="${WORKSPACE_DIR}/.venv/bin/python"
23+
if [ ! -f "${VENV_PYTHON}" ]; then
24+
VENV_PYTHON="/home/jackyf_google_com/maxtext/.venv/bin/python"
25+
fi
26+
if [ ! -f "${VENV_PYTHON}" ]; then
27+
VENV_PYTHON="python3"
28+
fi
29+
30+
# QWIX path auto-detection
31+
QWIX_DIR="/home/jackyf_google_com/qwix"
32+
if [ ! -d "${QWIX_DIR}" ]; then
33+
QWIX_DIR="${WORKSPACE_DIR}/../qwix"
34+
fi
35+
36+
export PYTHONPATH="${QWIX_DIR}:${WORKSPACE_DIR}/src"
37+
38+
echo "=========================================================="
39+
# Changing tile size to 16 for all Llama 3.1 runs to ensure 128 / 16 = 8 tiles,
40+
# which is divisible by the 8-device FSDP mesh.
41+
echo "Starting Flax NNX QLoRA Llama 3.1 8B End-to-End Test Suite"
42+
echo "Workspace: ${WORKSPACE_DIR}"
43+
echo "Python: ${VENV_PYTHON}"
44+
echo "PYTHONPATH: ${PYTHONPATH}"
45+
echo "=========================================================="
46+
47+
LOG_DIR="${WORKSPACE_DIR}/maxtext_output/e2e_llama31_logs"
48+
mkdir -p "${LOG_DIR}"
49+
50+
# Clean up old run artifacts
51+
echo "Cleaning up previous output directories..."
52+
rm -rf "${WORKSPACE_DIR}/maxtext_output/test_llama31_8b_scantrue_tile16"
53+
rm -rf "${WORKSPACE_DIR}/maxtext_output/test_llama31_8b_scanfalse_tile16"
54+
rm -rf "${WORKSPACE_DIR}/maxtext_output/test_llama31_8b_scantrue_sft_native_tile16"
55+
rm -rf "${WORKSPACE_DIR}/maxtext_output/test_llama31_8b_scanfalse_sft_native_tile16"
56+
rm -rf "${WORKSPACE_DIR}/maxtext_output/test_llama31_8b_scantrue_sft_custom_tile16"
57+
rm -rf "${WORKSPACE_DIR}/maxtext_output/test_llama31_8b_scanfalse_sft_custom_tile16"
58+
59+
# 1. RUN PRE-TRAIN BENCHMARKS
60+
echo -e "\n[1/6] Running Llama 3.1 Pre-train QLoRA (scan_layers=True) Benchmark..."
61+
"${VENV_PYTHON}" -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
62+
run_name=test_llama31_8b_scantrue_tile16 \
63+
model_name=llama3.1-8b scan_layers=True pure_nnx=True dataset_type=synthetic steps=20 \
64+
enable_checkpointing=False enable_goodput_recording=False enable_checkpoint_cloud_logger=False monitor_goodput=False \
65+
override_model_config=True base_num_decoder_layers=8 base_emb_dim=128 base_mlp_dim=256 base_num_query_heads=4 base_num_kv_heads=4 head_dim=128 \
66+
max_target_length=128 vocab_size=256 per_device_batch_size=1 \
67+
lora.enable_lora=True lora.lora_rank=4 lora.lora_alpha=8.0 lora.lora_weight_qtype=int8 lora.lora_tile_size=16 sharding_tolerance=1.0 \
68+
> "${LOG_DIR}/pre_train_scantrue.log" 2>&1
69+
PRE_TRAIN_SCANTRUE_STATUS=$?
70+
echo "Pre-train (scan_layers=True) Exit Status: ${PRE_TRAIN_SCANTRUE_STATUS}"
71+
72+
echo -e "\n[2/6] Running Llama 3.1 Pre-train QLoRA (scan_layers=False) Benchmark..."
73+
"${VENV_PYTHON}" -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
74+
run_name=test_llama31_8b_scanfalse_tile16 \
75+
model_name=llama3.1-8b scan_layers=False pure_nnx=True dataset_type=synthetic steps=20 \
76+
enable_checkpointing=False enable_goodput_recording=False enable_checkpoint_cloud_logger=False monitor_goodput=False \
77+
override_model_config=True base_num_decoder_layers=8 base_emb_dim=128 base_mlp_dim=256 base_num_query_heads=4 base_num_kv_heads=4 head_dim=128 \
78+
max_target_length=128 vocab_size=256 per_device_batch_size=1 \
79+
lora.enable_lora=True lora.lora_rank=4 lora.lora_alpha=8.0 lora.lora_weight_qtype=int8 lora.lora_tile_size=16 sharding_tolerance=1.0 \
80+
> "${LOG_DIR}/pre_train_scanfalse.log" 2>&1
81+
PRE_TRAIN_SCANFALSE_STATUS=$?
82+
echo "Pre-train (scan_layers=False) Exit Status: ${PRE_TRAIN_SCANFALSE_STATUS}"
83+
84+
# 2. RUN SFT-NATIVE BENCHMARKS
85+
echo -e "\n[3/6] Running Llama 3.1 SFT-Native QLoRA (scan_layers=True) Benchmark..."
86+
"${VENV_PYTHON}" -m maxtext.trainers.post_train.sft.train_sft_native src/maxtext/configs/base.yml \
87+
run_name=test_llama31_8b_scantrue_sft_native_tile16 \
88+
model_name=llama3.1-8b scan_layers=True pure_nnx=True dataset_type=synthetic steps=20 \
89+
enable_checkpointing=False enable_goodput_recording=False enable_checkpoint_cloud_logger=False monitor_goodput=False \
90+
override_model_config=True base_num_decoder_layers=8 base_emb_dim=128 base_mlp_dim=256 base_num_query_heads=4 base_num_kv_heads=4 head_dim=128 \
91+
max_target_length=128 vocab_size=256 per_device_batch_size=1 \
92+
lora.enable_lora=True lora.lora_rank=4 lora.lora_alpha=8.0 lora.lora_weight_qtype=int8 lora.lora_tile_size=16 sharding_tolerance=1.0 \
93+
> "${LOG_DIR}/sft_native_scantrue.log" 2>&1
94+
SFT_NATIVE_SCANTRUE_STATUS=$?
95+
echo "SFT-Native (scan_layers=True) Exit Status: ${SFT_NATIVE_SCANTRUE_STATUS}"
96+
97+
echo -e "\n[4/6] Running Llama 3.1 SFT-Native QLoRA (scan_layers=False) Benchmark..."
98+
"${VENV_PYTHON}" -m maxtext.trainers.post_train.sft.train_sft_native src/maxtext/configs/base.yml \
99+
run_name=test_llama31_8b_scanfalse_sft_native_tile16 \
100+
model_name=llama3.1-8b scan_layers=False pure_nnx=True dataset_type=synthetic steps=20 \
101+
enable_checkpointing=False enable_goodput_recording=False enable_checkpoint_cloud_logger=False monitor_goodput=False \
102+
override_model_config=True base_num_decoder_layers=8 base_emb_dim=128 base_mlp_dim=256 base_num_query_heads=4 base_num_kv_heads=4 head_dim=128 \
103+
max_target_length=128 vocab_size=256 per_device_batch_size=1 \
104+
lora.enable_lora=True lora.lora_rank=4 lora.lora_alpha=8.0 lora.lora_weight_qtype=int8 lora.lora_tile_size=16 sharding_tolerance=1.0 \
105+
> "${LOG_DIR}/sft_native_scanfalse.log" 2>&1
106+
SFT_NATIVE_SCANFALSE_STATUS=$?
107+
echo "SFT-Native (scan_layers=False) Exit Status: ${SFT_NATIVE_SCANFALSE_STATUS}"
108+
109+
# 3. RUN SFT-CUSTOM BENCHMARKS
110+
echo -e "\n[5/6] Running Llama 3.1 SFT-Custom QLoRA (scan_layers=True) Benchmark..."
111+
"${VENV_PYTHON}" -m maxtext.trainers.post_train.sft.train_sft src/maxtext/configs/post_train/sft.yml \
112+
run_name=test_llama31_8b_scantrue_sft_custom_tile16 \
113+
model_name=llama3.1-8b scan_layers=True pure_nnx=True dataset_type=synthetic steps=20 \
114+
enable_checkpointing=False enable_goodput_recording=False enable_checkpoint_cloud_logger=False monitor_goodput=False \
115+
override_model_config=True base_num_decoder_layers=8 base_emb_dim=128 base_mlp_dim=256 base_num_query_heads=4 base_num_kv_heads=4 head_dim=128 \
116+
max_target_length=128 vocab_size=256 per_device_batch_size=1 \
117+
lora.enable_lora=True lora.lora_rank=4 lora.lora_alpha=8.0 lora.lora_weight_qtype=int8 lora.lora_tile_size=16 sharding_tolerance=1.0 \
118+
> "${LOG_DIR}/sft_custom_scantrue.log" 2>&1
119+
SFT_CUSTOM_SCANTRUE_STATUS=$?
120+
echo "SFT-Custom (scan_layers=True) Exit Status: ${SFT_CUSTOM_SCANTRUE_STATUS}"
121+
122+
echo -e "\n[6/6] Running Llama 3.1 SFT-Custom QLoRA (scan_layers=False) Benchmark..."
123+
"${VENV_PYTHON}" -m maxtext.trainers.post_train.sft.train_sft src/maxtext/configs/post_train/sft.yml \
124+
run_name=test_llama31_8b_scanfalse_sft_custom_tile16 \
125+
model_name=llama3.1-8b scan_layers=False pure_nnx=True dataset_type=synthetic steps=20 \
126+
enable_checkpointing=False enable_goodput_recording=False enable_checkpoint_cloud_logger=False monitor_goodput=False \
127+
override_model_config=True base_num_decoder_layers=8 base_emb_dim=128 base_mlp_dim=256 base_num_query_heads=4 base_num_kv_heads=4 head_dim=128 \
128+
max_target_length=128 vocab_size=256 per_device_batch_size=1 \
129+
lora.enable_lora=True lora.lora_rank=4 lora.lora_alpha=8.0 lora.lora_weight_qtype=int8 lora.lora_tile_size=16 sharding_tolerance=1.0 \
130+
> "${LOG_DIR}/sft_custom_scanfalse.log" 2>&1
131+
SFT_CUSTOM_SCANFALSE_STATUS=$?
132+
echo "SFT-Custom (scan_layers=False) Exit Status: ${SFT_CUSTOM_SCANFALSE_STATUS}"
133+
134+
echo "=========================================================="
135+
echo "Parsing Performance Logs and Asserting Correctness"
136+
echo "=========================================================="
137+
138+
"${VENV_PYTHON}" - \
139+
"${LOG_DIR}/pre_train_scantrue.log" "${LOG_DIR}/pre_train_scanfalse.log" \
140+
"${LOG_DIR}/sft_native_scantrue.log" "${LOG_DIR}/sft_native_scanfalse.log" \
141+
"${LOG_DIR}/sft_custom_scantrue.log" "${LOG_DIR}/sft_custom_scanfalse.log" \
142+
"$PRE_TRAIN_SCANTRUE_STATUS" "$PRE_TRAIN_SCANFALSE_STATUS" \
143+
"$SFT_NATIVE_SCANTRUE_STATUS" "$SFT_NATIVE_SCANFALSE_STATUS" \
144+
"$SFT_CUSTOM_SCANTRUE_STATUS" "$SFT_CUSTOM_SCANFALSE_STATUS" << 'EOF'
145+
import sys
146+
import re
147+
import math
148+
149+
def parse_log(filepath):
150+
step_metrics = {}
151+
try:
152+
with open(filepath, "r") as f:
153+
for line in f:
154+
match = re.search(r"completed step:\s*(\d+),\s*seconds:\s*([\d\.]+),.*Tokens/s/device:\s*([\d\.]+),.*,\s*loss:\s*([\d\.]+)", line)
155+
if match:
156+
step = int(match.group(1))
157+
step_metrics[step] = {
158+
"time": float(match.group(2)),
159+
"tokens_s": float(match.group(3)),
160+
"loss": float(match.group(4))
161+
}
162+
except Exception as e:
163+
print(f"Warning: failed to read {filepath}: {e}")
164+
return step_metrics
165+
166+
pre_scantrue_log = sys.argv[1]
167+
pre_scanfalse_log = sys.argv[2]
168+
native_scantrue_log = sys.argv[3]
169+
native_scanfalse_log = sys.argv[4]
170+
custom_scantrue_log = sys.argv[5]
171+
custom_scanfalse_log = sys.argv[6]
172+
173+
m_pre_true = parse_log(pre_scantrue_log)
174+
m_pre_false = parse_log(pre_scanfalse_log)
175+
m_native_true = parse_log(native_scantrue_log)
176+
m_native_false = parse_log(native_scanfalse_log)
177+
m_custom_true = parse_log(custom_scantrue_log)
178+
m_custom_false = parse_log(custom_scanfalse_log)
179+
180+
print("\n### E2E Llama 3.1 QLoRA Benchmark Results (Steps 10-15)\n")
181+
print("| Trainer | Scan Layers | Step | Loss | Step Time | Tokens/s/dev |")
182+
print("|---|---|---|---|---|---|")
183+
184+
def format_steps(name, scan_str, metrics, steps_range):
185+
for step in steps_range:
186+
if step in metrics:
187+
loss_val = f"{metrics[step]['loss']:.4f}"
188+
t_val = f"{metrics[step]['time']:.3f}s"
189+
tok_val = f"{metrics[step]['tokens_s']:.1f}"
190+
print(f"| {name} | {scan_str} | {step} | {loss_val} | {t_val} | {tok_val} |")
191+
else:
192+
print(f"| {name} | {scan_str} | {step} | N/A | N/A | N/A |")
193+
194+
format_steps("Pre-train", "True", m_pre_true, range(10, 16))
195+
format_steps("Pre-train", "False", m_pre_false, range(10, 16))
196+
format_steps("SFT-Native", "True", m_native_true, range(10, 16))
197+
format_steps("SFT-Native", "False", m_native_false, range(10, 16))
198+
format_steps("SFT-Custom", "True", m_custom_true, range(10, 16))
199+
format_steps("SFT-Custom", "False", m_custom_false, range(10, 16))
200+
201+
print("\nValidating results...")
202+
203+
success = True
204+
configs = [
205+
("Pre-train Scan True", m_pre_true, int(sys.argv[7])),
206+
("Pre-train Scan False", m_pre_false, int(sys.argv[8])),
207+
("SFT-Native Scan True", m_native_true, int(sys.argv[9])),
208+
("SFT-Native Scan False", m_native_false, int(sys.argv[10])),
209+
("SFT-Custom Scan True", m_custom_true, int(sys.argv[11])),
210+
("SFT-Custom Scan False", m_custom_false, int(sys.argv[12])),
211+
]
212+
213+
for name, metrics, exit_status in configs:
214+
if exit_status != 0:
215+
print(f"Error: {name} process exited with failure code {exit_status}!")
216+
success = False
217+
if not metrics:
218+
print(f"Error: {name} benchmark did not output any completed step log metrics!")
219+
success = False
220+
continue
221+
222+
for step in range(10, 16):
223+
if step not in metrics:
224+
print(f"Error: {name} is missing metrics for step {step}")
225+
success = False
226+
continue
227+
228+
loss = metrics[step]["loss"]
229+
if math.isnan(loss) or math.isinf(loss) or loss < 1.0 or loss > 15.0:
230+
print(f"Error: {name} step {step} has anomalous loss: {loss}")
231+
success = False
232+
233+
if success:
234+
print("Verification: ALL CHECKS PASSED SUCCESSFULLY!")
235+
sys.exit(0)
236+
else:
237+
print("Verification: SOME CHECKS FAILED.")
238+
sys.exit(1)
239+
EOF
240+
241+
E2E_STATUS=$?
242+
exit $E2E_STATUS

0 commit comments

Comments
 (0)