Keep the full LoRA delta when merging into bfloat16 weights

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,

src/maxdiffusion/tests/lora_nnx_test.py

@@ -0,0 +1,132 @@
+"""
+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
+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.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):
+  """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.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.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.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)
+      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.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)
+    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.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.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.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}
+
+    # 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)
+
+
+if __name__ == "__main__":
+  unittest.main()