From 2313406a0e32994464540976d29b093e874cee20 Mon Sep 17 00:00:00 2001 From: boshjerns <98094229+boshjerns@users.noreply.github.com> Date: Sun, 18 Oct 2026 18:52:04 +0200 Subject: [PATCH 1/3] Keep the full LoRA delta when merging into bfloat16 weights --- src/maxdiffusion/models/lora_nnx.py | 60 ++++++++--- src/maxdiffusion/tests/lora_nnx_test.py | 129 ++++++++++++++++++++++++ 2 files changed, 176 insertions(+), 13 deletions(-) create mode 100644 src/maxdiffusion/tests/lora_nnx_test.py diff --git a/src/maxdiffusion/models/lora_nnx.py b/src/maxdiffusion/models/lora_nnx.py index 97456630b..4d30b9b22 100644 --- a/src/maxdiffusion/models/lora_nnx.py +++ b/src/maxdiffusion/models/lora_nnx.py @@ -28,51 +28,77 @@ # ----------------------------------------------------------------------------- +def _add_and_round(param, delta, key): + """ + Returns param + delta, summed in float32 and rounded to param.dtype once. + A LoRA delta is often smaller than half the bfloat16 spacing of the weight, so + round-to-nearest would drop most of it. bfloat16 results are therefore rounded + stochastically, which keeps the full delta in expectation. + """ + total = param.astype(jnp.float32) + delta.astype(jnp.float32) + if param.dtype != jnp.bfloat16: + return total.astype(param.dtype) + # Add random bits below the bfloat16 mantissa, then truncate to the upper 16 bits. + bits = jax.lax.bitcast_convert_type(total, jnp.uint32) + noise = jax.random.bits(key, total.shape, jnp.uint32) & jnp.uint32(0xFFFF) + rounded = jax.lax.bitcast_convert_type((bits + noise) & jnp.uint32(0xFFFF0000), jnp.float32) + return rounded.astype(jnp.bfloat16) + + @jax.jit -def _compute_and_add_single_jit(kernel, bias, down, up, scale, w_diff, b_diff): +def _compute_and_add_single_jit(kernel, bias, down, up, scale, w_diff, b_diff, key): """ Applies LoRA + Weight Diff + Bias Diff on device. """ + kernel_key, bias_key = jax.random.split(key) + kernel_delta = jnp.zeros(kernel.shape, jnp.float32) + # 1. Apply LoRA (if valid) if down is not None and up is not None: # down: (Rank, In), up: (Out, Rank) -> Result: (In, Out) # Note: We reshape to kernel shape to handle 1x1 convs delta = (down.T @ up.T).reshape(kernel.shape) - kernel = kernel + (delta * scale).astype(kernel.dtype) + kernel_delta = kernel_delta + delta * scale # 2. Apply Full Weight Diff (if valid) if w_diff is not None: - kernel = kernel + w_diff.astype(kernel.dtype) + kernel_delta = kernel_delta + w_diff + + kernel = _add_and_round(kernel, kernel_delta, kernel_key) # 3. Apply Bias Diff (if valid and bias exists) if bias is not None and b_diff is not None: - bias = bias + b_diff.astype(bias.dtype) + bias = _add_and_round(bias, b_diff, bias_key) return kernel, bias @jax.jit -def _compute_and_add_scanned_jit(kernel, downs, ups, alphas, global_scale, w_diffs=None, b_diffs=None, bias=None): +def _compute_and_add_scanned_jit(kernel, downs, ups, alphas, global_scale, key, w_diffs=None, b_diffs=None, bias=None): """ Applies scanned LoRA + Diffs. """ + kernel_key, bias_key = jax.random.split(key) + kernel_delta = jnp.zeros(kernel.shape, jnp.float32) + # 1. Apply LoRA if downs is not None and ups is not None: rank = downs.shape[1] scales = global_scale * alphas / rank # Batch Matmul: (L, In, Out) delta = jnp.matmul(jnp.swapaxes(downs, 1, 2), jnp.swapaxes(ups, 1, 2)) - delta = (delta * scales).astype(kernel.dtype) - kernel = kernel + delta.reshape(kernel.shape) + kernel_delta = kernel_delta + (delta * scales).reshape(kernel.shape) # 2. Apply Scanned Weight Diffs (L, ...) if w_diffs is not None: - kernel = kernel + w_diffs.astype(kernel.dtype) + kernel_delta = kernel_delta + w_diffs + + kernel = _add_and_round(kernel, kernel_delta, kernel_key) # 3. Apply Scanned Bias Diffs (L, ...) # Note: Scanned bias is usually shape (L, Out) if bias is not None and b_diffs is not None: - bias = bias + b_diffs.astype(bias.dtype) + bias = _add_and_round(bias, b_diffs, bias_key) return kernel, bias @@ -147,7 +173,7 @@ def parse_lora_dict(state_dict, dtype): return lora_params -def _merge_lora_layer(module, weights, scale): +def _merge_lora_layer(module, weights, scale, key): """Merges LoRA weights into a single non-scanned layer.""" is_conv_kxk_locon = False if isinstance(module, nnx.Conv) and module.kernel_size != (1, 1) and "down" in weights and "up" in weights: @@ -232,7 +258,7 @@ def _merge_lora_layer(module, weights, scale): # --- EXECUTE JIT UPDATE --- if down_w is not None or w_diff is not None or b_diff is not None: new_kernel, new_bias = _compute_and_add_single_jit( - module.kernel.value, bias_val, down_w, up_w, current_scale, w_diff, b_diff + module.kernel.value, bias_val, down_w, up_w, current_scale, w_diff, b_diff, key ) module.kernel.value = new_kernel @@ -253,6 +279,8 @@ def merge_lora(model: nnx.Module, state_dict: dict, rank: int, scale: float, tra lora_params = parse_lora_dict(state_dict, dtype=dtype) max_logging.log(f"Parsed {len(lora_params)} unique module keys.") matched_keys = set() + # Fixed key so stochastic rounding of the merged weights is reproducible. + key = jax.random.key(0) assigned_count = 0 for path, module in nnx.iter_graph(model): @@ -265,7 +293,8 @@ def merge_lora(model: nnx.Module, state_dict: dict, rank: int, scale: float, tra if lora_key and lora_key in lora_params: matched_keys.add(lora_key) weights = lora_params[lora_key] - if _merge_lora_layer(module, weights, scale): + key, layer_key = jax.random.split(key) + if _merge_lora_layer(module, weights, scale, layer_key): assigned_count += 1 max_logging.log(f"Merged weights into {assigned_count} layers.") @@ -286,6 +315,8 @@ def merge_lora_for_scanned( lora_params = parse_lora_dict(state_dict, dtype=dtype) max_logging.log(f"Parsed {len(lora_params)} keys for scanned merge.") matched_keys = set() + # Fixed key so stochastic rounding of the merged weights is reproducible. + key = jax.random.key(0) assigned_count = 0 for path, module in nnx.iter_graph(model): @@ -317,7 +348,8 @@ def merge_lora_for_scanned( if lora_key in lora_params: matched_keys.add(lora_key) weights = lora_params[lora_key] - if _merge_lora_layer(module, weights, scale): + key, layer_key = jax.random.split(key) + if _merge_lora_layer(module, weights, scale, layer_key): assigned_count += 1 continue @@ -476,12 +508,14 @@ def merge_lora_for_scanned( bias_val = module.bias.value if module.bias is not None else None # Call JIT + key, layer_key = jax.random.split(key) new_k, new_b = _compute_and_add_scanned_jit( module.kernel.value, stack_down if has_lora else None, stack_up if has_lora else None, stack_alpha if has_lora else None, scale, + layer_key, stack_w_diff, stack_b_diff, bias_val, diff --git a/src/maxdiffusion/tests/lora_nnx_test.py b/src/maxdiffusion/tests/lora_nnx_test.py new file mode 100644 index 000000000..2ddd48e4b --- /dev/null +++ b/src/maxdiffusion/tests/lora_nnx_test.py @@ -0,0 +1,129 @@ +""" +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. +""" + +import unittest + +import jax.numpy as jnp +import numpy as np +from flax import nnx + +from maxdiffusion.models import lora_nnx + +IN_FEATURES, OUT_FEATURES, RANK, NUM_LAYERS = 64, 4096, 8, 4 + + +class _Model(nnx.Module): + + def __init__(self, dtype, num_layers=None): + self.proj = nnx.Linear(IN_FEATURES, OUT_FEATURES, param_dtype=dtype, rngs=nnx.Rngs(0)) + if num_layers is not None: + # Stacked weights, as in a scanned transformer. + self.proj.kernel.set_value(jnp.zeros((num_layers, IN_FEATURES, OUT_FEATURES), dtype)) + self.proj.bias.set_value(jnp.zeros((num_layers, OUT_FEATURES), dtype)) + + +def _lora(rng, base_std, delta_ratio): + """Returns (lora_A, lora_B, delta) with delta = (B @ A).T scaled to delta_ratio * base_std.""" + down = rng.normal(size=(RANK, IN_FEATURES)).astype(np.float32) + up = rng.normal(size=(OUT_FEATURES, RANK)).astype(np.float32) + delta = (up @ down).T + factor = np.sqrt(delta_ratio * base_std / delta.std()) + return down * factor, up * factor, delta * factor**2 + + +def _retained(merged, base, delta): + """Fraction of delta present in merged - base (1.0 means the full delta was applied).""" + diff = np.asarray(merged, np.float32) - np.asarray(base, np.float32) + return float((diff * delta).sum() / (delta * delta).sum()) + + +class LoraNNXMergeTest(unittest.TestCase): + """Tests merging LoRA weights into NNX layers.""" + + # LoRA updates are typically far smaller than the weights they modify. At this ratio + # round-to-nearest into bfloat16 keeps only about half of the update. + BASE_STD, DELTA_RATIO = 0.02, 1e-3 + + def _base(self, rng, shape): + return rng.normal(0, self.BASE_STD, shape).astype(np.float32) + + def test_merge_keeps_small_delta_in_bfloat16(self): + rng = np.random.default_rng(0) + model = _Model(jnp.bfloat16) + model.proj.kernel.set_value(jnp.asarray(self._base(rng, (IN_FEATURES, OUT_FEATURES)), jnp.bfloat16)) + model.proj.bias.set_value(jnp.asarray(self._base(rng, (OUT_FEATURES,)), jnp.bfloat16)) + kernel, bias = np.asarray(model.proj.kernel[...]), np.asarray(model.proj.bias[...]) + down, up, delta = _lora(rng, self.BASE_STD, self.DELTA_RATIO) + bias_delta = rng.normal(0, self.BASE_STD * self.DELTA_RATIO, OUT_FEATURES).astype(np.float32) + state_dict = {"proj.lora_A.weight": down, "proj.lora_B.weight": up, "proj.diff_b": bias_delta} + + lora_nnx.merge_lora(model, state_dict, RANK, 1.0, lambda path: "proj" if path == "proj" else None, dtype="bfloat16") + + self.assertEqual(model.proj.kernel[...].dtype, jnp.bfloat16) + self.assertAlmostEqual(_retained(model.proj.kernel[...], kernel, delta), 1.0, delta=0.05) + self.assertAlmostEqual(_retained(model.proj.bias[...], bias, bias_delta), 1.0, delta=0.1) + + def test_scanned_merge_keeps_small_delta_in_bfloat16(self): + rng = np.random.default_rng(1) + model = _Model(jnp.bfloat16, num_layers=NUM_LAYERS) + model.proj.kernel.set_value(jnp.asarray(self._base(rng, (NUM_LAYERS, IN_FEATURES, OUT_FEATURES)), jnp.bfloat16)) + model.proj.bias.set_value(jnp.asarray(self._base(rng, (NUM_LAYERS, OUT_FEATURES)), jnp.bfloat16)) + kernel, bias = np.asarray(model.proj.kernel[...]), np.asarray(model.proj.bias[...]) + state_dict, deltas, bias_deltas = {}, [], [] + for i in range(NUM_LAYERS): + down, up, delta = _lora(rng, self.BASE_STD, self.DELTA_RATIO) + bias_delta = rng.normal(0, self.BASE_STD * self.DELTA_RATIO, OUT_FEATURES).astype(np.float32) + state_dict.update({f"proj.{i}.lora_A.weight": down, f"proj.{i}.lora_B.weight": up, f"proj.{i}.diff_b": bias_delta}) + deltas.append(delta) + bias_deltas.append(bias_delta) + + lora_nnx.merge_lora_for_scanned( + model, state_dict, RANK, 1.0, lambda path: "proj.{}" if path == "proj" else None, dtype="bfloat16" + ) + + self.assertEqual(model.proj.kernel[...].dtype, jnp.bfloat16) + self.assertAlmostEqual(_retained(model.proj.kernel[...], kernel, np.stack(deltas)), 1.0, delta=0.05) + self.assertAlmostEqual(_retained(model.proj.bias[...], bias, np.stack(bias_deltas)), 1.0, delta=0.1) + + def test_merge_is_deterministic(self): + rng = np.random.default_rng(2) + base = jnp.asarray(self._base(rng, (IN_FEATURES, OUT_FEATURES)), jnp.bfloat16) + down, up, _ = _lora(rng, self.BASE_STD, self.DELTA_RATIO) + state_dict = {"proj.lora_A.weight": down, "proj.lora_B.weight": up} + merged = [] + for _ in range(2): + model = _Model(jnp.bfloat16) + model.proj.kernel.set_value(base) + lora_nnx.merge_lora(model, state_dict, RANK, 1.0, lambda path: "proj" if path == "proj" else None) + merged.append(np.asarray(model.proj.kernel[...], np.float32)) + np.testing.assert_array_equal(merged[0], merged[1]) + + def test_merge_float32_is_exact_sum(self): + rng = np.random.default_rng(3) + model = _Model(jnp.float32) + kernel = self._base(rng, (IN_FEATURES, OUT_FEATURES)) + model.proj.kernel.set_value(jnp.asarray(kernel)) + down, up, delta = _lora(rng, self.BASE_STD, self.DELTA_RATIO) + state_dict = {"proj.lora_A.weight": down, "proj.lora_B.weight": up} + + lora_nnx.merge_lora(model, state_dict, RANK, 1.0, lambda path: "proj" if path == "proj" else None) + + self.assertEqual(model.proj.kernel[...].dtype, jnp.float32) + np.testing.assert_allclose(np.asarray(model.proj.kernel[...]), kernel + delta, rtol=1e-6, atol=1e-7) + + +if __name__ == "__main__": + unittest.main() From 6215b48a7f9e0597121db0529a2690c17c1f0474 Mon Sep 17 00:00:00 2001 From: boshjerns <98094229+boshjerns@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:27:19 +0200 Subject: [PATCH 2/3] Use .value for flax variables in the LoRA merge test --- src/maxdiffusion/tests/lora_nnx_test.py | 38 ++++++++++++------------- 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/src/maxdiffusion/tests/lora_nnx_test.py b/src/maxdiffusion/tests/lora_nnx_test.py index 2ddd48e4b..bca481dab 100644 --- a/src/maxdiffusion/tests/lora_nnx_test.py +++ b/src/maxdiffusion/tests/lora_nnx_test.py @@ -31,8 +31,8 @@ def __init__(self, dtype, num_layers=None): self.proj = nnx.Linear(IN_FEATURES, OUT_FEATURES, param_dtype=dtype, rngs=nnx.Rngs(0)) if num_layers is not None: # Stacked weights, as in a scanned transformer. - self.proj.kernel.set_value(jnp.zeros((num_layers, IN_FEATURES, OUT_FEATURES), dtype)) - self.proj.bias.set_value(jnp.zeros((num_layers, OUT_FEATURES), dtype)) + self.proj.kernel.value = jnp.zeros((num_layers, IN_FEATURES, OUT_FEATURES), dtype) + self.proj.bias.value = jnp.zeros((num_layers, OUT_FEATURES), dtype) def _lora(rng, base_std, delta_ratio): @@ -63,25 +63,25 @@ def _base(self, rng, shape): def test_merge_keeps_small_delta_in_bfloat16(self): rng = np.random.default_rng(0) model = _Model(jnp.bfloat16) - model.proj.kernel.set_value(jnp.asarray(self._base(rng, (IN_FEATURES, OUT_FEATURES)), jnp.bfloat16)) - model.proj.bias.set_value(jnp.asarray(self._base(rng, (OUT_FEATURES,)), jnp.bfloat16)) - kernel, bias = np.asarray(model.proj.kernel[...]), np.asarray(model.proj.bias[...]) + model.proj.kernel.value = jnp.asarray(self._base(rng, (IN_FEATURES, OUT_FEATURES)), jnp.bfloat16) + model.proj.bias.value = jnp.asarray(self._base(rng, (OUT_FEATURES,)), jnp.bfloat16) + kernel, bias = np.asarray(model.proj.kernel.value), np.asarray(model.proj.bias.value) down, up, delta = _lora(rng, self.BASE_STD, self.DELTA_RATIO) bias_delta = rng.normal(0, self.BASE_STD * self.DELTA_RATIO, OUT_FEATURES).astype(np.float32) state_dict = {"proj.lora_A.weight": down, "proj.lora_B.weight": up, "proj.diff_b": bias_delta} lora_nnx.merge_lora(model, state_dict, RANK, 1.0, lambda path: "proj" if path == "proj" else None, dtype="bfloat16") - self.assertEqual(model.proj.kernel[...].dtype, jnp.bfloat16) - self.assertAlmostEqual(_retained(model.proj.kernel[...], kernel, delta), 1.0, delta=0.05) - self.assertAlmostEqual(_retained(model.proj.bias[...], bias, bias_delta), 1.0, delta=0.1) + self.assertEqual(model.proj.kernel.value.dtype, jnp.bfloat16) + self.assertAlmostEqual(_retained(model.proj.kernel.value, kernel, delta), 1.0, delta=0.05) + self.assertAlmostEqual(_retained(model.proj.bias.value, bias, bias_delta), 1.0, delta=0.1) def test_scanned_merge_keeps_small_delta_in_bfloat16(self): rng = np.random.default_rng(1) model = _Model(jnp.bfloat16, num_layers=NUM_LAYERS) - model.proj.kernel.set_value(jnp.asarray(self._base(rng, (NUM_LAYERS, IN_FEATURES, OUT_FEATURES)), jnp.bfloat16)) - model.proj.bias.set_value(jnp.asarray(self._base(rng, (NUM_LAYERS, OUT_FEATURES)), jnp.bfloat16)) - kernel, bias = np.asarray(model.proj.kernel[...]), np.asarray(model.proj.bias[...]) + model.proj.kernel.value = jnp.asarray(self._base(rng, (NUM_LAYERS, IN_FEATURES, OUT_FEATURES)), jnp.bfloat16) + model.proj.bias.value = jnp.asarray(self._base(rng, (NUM_LAYERS, OUT_FEATURES)), jnp.bfloat16) + kernel, bias = np.asarray(model.proj.kernel.value), np.asarray(model.proj.bias.value) state_dict, deltas, bias_deltas = {}, [], [] for i in range(NUM_LAYERS): down, up, delta = _lora(rng, self.BASE_STD, self.DELTA_RATIO) @@ -94,9 +94,9 @@ def test_scanned_merge_keeps_small_delta_in_bfloat16(self): model, state_dict, RANK, 1.0, lambda path: "proj.{}" if path == "proj" else None, dtype="bfloat16" ) - self.assertEqual(model.proj.kernel[...].dtype, jnp.bfloat16) - self.assertAlmostEqual(_retained(model.proj.kernel[...], kernel, np.stack(deltas)), 1.0, delta=0.05) - self.assertAlmostEqual(_retained(model.proj.bias[...], bias, np.stack(bias_deltas)), 1.0, delta=0.1) + self.assertEqual(model.proj.kernel.value.dtype, jnp.bfloat16) + self.assertAlmostEqual(_retained(model.proj.kernel.value, kernel, np.stack(deltas)), 1.0, delta=0.05) + self.assertAlmostEqual(_retained(model.proj.bias.value, bias, np.stack(bias_deltas)), 1.0, delta=0.1) def test_merge_is_deterministic(self): rng = np.random.default_rng(2) @@ -106,23 +106,23 @@ def test_merge_is_deterministic(self): merged = [] for _ in range(2): model = _Model(jnp.bfloat16) - model.proj.kernel.set_value(base) + model.proj.kernel.value = base lora_nnx.merge_lora(model, state_dict, RANK, 1.0, lambda path: "proj" if path == "proj" else None) - merged.append(np.asarray(model.proj.kernel[...], np.float32)) + merged.append(np.asarray(model.proj.kernel.value, np.float32)) np.testing.assert_array_equal(merged[0], merged[1]) def test_merge_float32_is_exact_sum(self): rng = np.random.default_rng(3) model = _Model(jnp.float32) kernel = self._base(rng, (IN_FEATURES, OUT_FEATURES)) - model.proj.kernel.set_value(jnp.asarray(kernel)) + model.proj.kernel.value = jnp.asarray(kernel) down, up, delta = _lora(rng, self.BASE_STD, self.DELTA_RATIO) state_dict = {"proj.lora_A.weight": down, "proj.lora_B.weight": up} lora_nnx.merge_lora(model, state_dict, RANK, 1.0, lambda path: "proj" if path == "proj" else None) - self.assertEqual(model.proj.kernel[...].dtype, jnp.float32) - np.testing.assert_allclose(np.asarray(model.proj.kernel[...]), kernel + delta, rtol=1e-6, atol=1e-7) + self.assertEqual(model.proj.kernel.value.dtype, jnp.float32) + np.testing.assert_allclose(np.asarray(model.proj.kernel.value), kernel + delta, rtol=1e-6, atol=1e-7) if __name__ == "__main__": From 4b3d8cb4ed870c78ab5b0a0aad53429ec5257889 Mon Sep 17 00:00:00 2001 From: boshjerns <98094229+boshjerns@users.noreply.github.com> Date: Mon, 19 Oct 2026 18:31:39 +0200 Subject: [PATCH 3/3] Check the float32 LoRA merge at full matmul precision --- src/maxdiffusion/tests/lora_nnx_test.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/maxdiffusion/tests/lora_nnx_test.py b/src/maxdiffusion/tests/lora_nnx_test.py index bca481dab..67c88a3f5 100644 --- a/src/maxdiffusion/tests/lora_nnx_test.py +++ b/src/maxdiffusion/tests/lora_nnx_test.py @@ -16,6 +16,7 @@ import unittest +import jax import jax.numpy as jnp import numpy as np from flax import nnx @@ -119,7 +120,9 @@ def test_merge_float32_is_exact_sum(self): down, up, delta = _lora(rng, self.BASE_STD, self.DELTA_RATIO) state_dict = {"proj.lora_A.weight": down, "proj.lora_B.weight": up} - lora_nnx.merge_lora(model, state_dict, RANK, 1.0, lambda path: "proj" if path == "proj" else None) + # TPUs run float32 matmuls in reduced precision by default; check the sum at full precision. + with jax.default_matmul_precision("highest"): + lora_nnx.merge_lora(model, state_dict, RANK, 1.0, lambda path: "proj" if path == "proj" else None) self.assertEqual(model.proj.kernel.value.dtype, jnp.float32) np.testing.assert_allclose(np.asarray(model.proj.kernel.value), kernel + delta, rtol=1e-6, atol=1e-7)