diff --git a/tests/diffusion/types_test.py b/tests/diffusion/types_test.py
new file mode 100644
index 000000000..16585fc1e
--- /dev/null
+++ b/tests/diffusion/types_test.py
@@ -0,0 +1,207 @@
+# 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.
+
+"""Tests for framework-neutral diffusion contracts."""
+
+from absl.testing import absltest
+from flax import nnx
+import jax
+import jax.numpy as jnp
+import jax.sharding as shd
+import numpy as np
+from tunix.diffusion import interfaces
+from tunix.diffusion import types
+from tunix.sft import sharding_utils
+
+
+def _batch(
+    *,
+    model_inputs=None,
+    target_ids=None,
+    loss_weights=None,
+) -> types.DiffusionTokenBatch:
+  if model_inputs is None:
+    model_inputs = {
+        "input_tokens": jnp.array([[1, 2, 3], [4, 5, 6]], dtype=jnp.int32),
+        "noise_levels": jnp.array([0.25, 0.75], dtype=jnp.float32),
+    }
+  if target_ids is None:
+    target_ids = jnp.array([[2, 3, 4], [5, 6, 7]], dtype=jnp.int32)
+  if loss_weights is None:
+    loss_weights = jnp.ones((2, 3), dtype=jnp.float32)
+  return types.DiffusionTokenBatch.create(
+      model_inputs=model_inputs,
+      target_ids=target_ids,
+      loss_weights=loss_weights,
+  )
+
+
+class _ToyModel(nnx.Module):
+
+  def __init__(self, vocab_size: int):
+    self.bias = nnx.Param(jnp.arange(vocab_size, dtype=jnp.float32))
+
+  def score(self, input_tokens: jax.Array) -> jax.Array:
+    bias = self.bias[...]
+    return jax.nn.one_hot(input_tokens, bias.shape[0]) + bias
+
+
+def _score_fn(model: nnx.Module, model_inputs: types.ModelInputs) -> jax.Array:
+  return model.score(model_inputs["input_tokens"])
+
+
+class DiffusionTokenBatchTest(absltest.TestCase):
+
+  def test_requires_rank_two_integer_targets(self):
+    with self.assertRaisesRegex(ValueError, "shape \\[batch, length\\]"):
+      _batch(target_ids=jnp.ones((2, 3, 1), dtype=jnp.int32))
+    with self.assertRaisesRegex(TypeError, "integer dtype"):
+      _batch(target_ids=jnp.ones((2, 3), dtype=jnp.float32))
+
+  def test_requires_matching_real_loss_weights(self):
+    with self.assertRaisesRegex(ValueError, "must match target_ids shape"):
+      _batch(loss_weights=jnp.ones((2, 2), dtype=jnp.float32))
+    with self.assertRaisesRegex(TypeError, "real numeric or boolean"):
+      _batch(loss_weights=jnp.ones((2, 3), dtype=jnp.complex64))
+
+  def test_accepts_boolean_and_integer_loss_weights(self):
+    self.assertEqual(
+        _batch(loss_weights=jnp.ones((2, 3), dtype=bool)).loss_weights.dtype,
+        jnp.bool_,
+    )
+    self.assertEqual(
+        _batch(
+            loss_weights=jnp.ones((2, 3), dtype=jnp.int32)
+        ).loss_weights.dtype,
+        jnp.int32,
+    )
+
+  def test_rejects_invalid_loss_weight_values(self):
+    with self.assertRaisesRegex(ValueError, "finite"):
+      _batch(loss_weights=jnp.array([[jnp.nan, 1, 1], [1, 1, 1]]))
+    with self.assertRaisesRegex(ValueError, "finite"):
+      _batch(loss_weights=jnp.array([[jnp.inf, 1, 1], [1, 1, 1]]))
+    with self.assertRaisesRegex(ValueError, "nonnegative"):
+      _batch(loss_weights=jnp.array([[-1.0, 1, 1], [1, 1, 1]]))
+
+  def test_rejects_negative_active_targets_but_allows_inactive_sentinels(self):
+    target_ids = jnp.array([[-1, 2, 3], [4, 5, 6]], dtype=jnp.int32)
+    with self.assertRaisesRegex(ValueError, "active target_ids"):
+      _batch(target_ids=target_ids)
+
+    batch = _batch(
+        target_ids=target_ids,
+        loss_weights=jnp.array([[0, 1, 1], [1, 1, 1]], dtype=jnp.float32),
+    )
+    self.assertEqual(int(batch.target_ids[0, 0]), -1)
+
+  def test_requires_nonempty_batch_major_array_inputs(self):
+    with self.assertRaisesRegex(ValueError, "at least one array"):
+      _batch(model_inputs={})
+    with self.assertRaisesRegex(TypeError, "must be a JAX or NumPy array"):
+      _batch(model_inputs={"tokens": [[1, 2, 3], [4, 5, 6]]})
+    with self.assertRaisesRegex(ValueError, "not scalar"):
+      _batch(model_inputs={"temperature": jnp.array(1.0)})
+    with self.assertRaisesRegex(ValueError, "batch size 3; expected 2"):
+      _batch(model_inputs={"tokens": jnp.ones((3, 3), dtype=jnp.int32)})
+    with self.assertRaisesRegex(TypeError, "real numeric or boolean"):
+      _batch(model_inputs={"tokens": np.array([["a"], ["b"]])})
+
+  def test_is_a_jax_pytree_and_jittable(self):
+    batch = _batch()
+    leaves, tree_def = jax.tree.flatten(batch)
+
+    self.assertLen(leaves, 4)
+    restored = jax.tree.unflatten(tree_def, leaves)
+    np.testing.assert_array_equal(restored.target_ids, batch.target_ids)
+
+    @jax.jit
+    def weighted_targets(value):
+      return value.target_ids * value.loss_weights
+
+    np.testing.assert_array_equal(
+        weighted_targets(batch), batch.target_ids.astype(jnp.float32)
+    )
+
+  def test_array_leaves_are_compatible_with_data_sharding(self):
+    batch = _batch()
+    mesh = shd.Mesh(np.array(jax.devices()[:1]), ("data",))
+
+    with mesh:
+      sharded_batch = sharding_utils.shard_input(batch, ("data",))
+
+    for leaf in jax.tree.leaves(sharded_batch):
+      self.assertIsInstance(leaf, jax.Array)
+      self.assertEqual(leaf.sharding.spec, shd.PartitionSpec("data"))
+
+
+class DiffusionLogitsTest(absltest.TestCase):
+
+  def test_scores_are_target_aligned_and_jittable(self):
+    batch = _batch()
+    model = _ToyModel(vocab_size=8)
+
+    @nnx.jit
+    def score(model, batch):
+      return interfaces.compute_diffusion_logits(model, batch, _score_fn)
+
+    scores = score(model, batch)
+    self.assertEqual(scores.shape, (2, 3, 8))
+    self.assertEqual(scores.dtype, jnp.float32)
+
+  def test_rejects_non_target_aligned_scores(self):
+    batch = _batch()
+    with self.assertRaisesRegex(ValueError, "align with target_ids"):
+      types.validate_diffusion_logits(
+          batch, jnp.ones((2, 2, 8), dtype=jnp.float32)
+      )
+    with self.assertRaisesRegex(ValueError, "non-empty vocabulary"):
+      types.validate_diffusion_logits(
+          batch, jnp.ones((2, 3, 0), dtype=jnp.float32)
+      )
+
+  def test_rejects_invalid_score_rank_and_dtype(self):
+    batch = _batch()
+    with self.assertRaisesRegex(ValueError, "shape \\[batch, length, vocab\\]"):
+      types.validate_diffusion_logits(
+          batch, jnp.ones((2, 3), dtype=jnp.float32)
+      )
+    with self.assertRaisesRegex(TypeError, "floating-point dtype"):
+      types.validate_diffusion_logits(
+          batch, jnp.ones((2, 3, 8), dtype=jnp.int32)
+      )
+
+  def test_rejects_active_targets_outside_vocabulary(self):
+    batch = _batch(
+        target_ids=jnp.array([[2, 8, 4], [5, 6, 7]], dtype=jnp.int32)
+    )
+    with self.assertRaisesRegex(ValueError, "smaller than vocabulary size 8"):
+      types.validate_diffusion_logits(
+          batch, jnp.ones((2, 3, 8), dtype=jnp.float32)
+      )
+
+  def test_validated_scorer_rejects_misaligned_implementation(self):
+    def misaligned_score_fn(model, model_inputs):
+      del model
+      batch_size = model_inputs["input_tokens"].shape[0]
+      return jnp.ones((batch_size, 1, 8), dtype=jnp.float32)
+
+    with self.assertRaisesRegex(ValueError, "align with target_ids"):
+      interfaces.compute_diffusion_logits(
+          _ToyModel(vocab_size=8), _batch(), misaligned_score_fn
+      )
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/distillation/diffusion_batch_test.py b/tests/distillation/diffusion_batch_test.py
new file mode 100644
index 000000000..18b3bb9b2
--- /dev/null
+++ b/tests/distillation/diffusion_batch_test.py
@@ -0,0 +1,67 @@
+# 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.
+
+"""Tests for external-teacher diffusion distillation batches."""
+
+from absl.testing import absltest
+import jax
+import jax.numpy as jnp
+import numpy as np
+from tunix.diffusion import types as diffusion_types
+from tunix.distillation import diffusion
+
+
+def _student_batch() -> diffusion_types.DiffusionTokenBatch:
+  return diffusion_types.DiffusionTokenBatch.create(
+      model_inputs={"example_ids": jnp.array([0], dtype=jnp.int32)},
+      target_ids=jnp.array([[0, 1]], dtype=jnp.int32),
+      loss_weights=jnp.ones((1, 2), dtype=jnp.float32),
+  )
+
+
+class DiffusionDistillationBatchTest(absltest.TestCase):
+
+  def test_rejects_non_diffusion_student_batch(self):
+    with self.assertRaisesRegex(TypeError, "must be a DiffusionTokenBatch"):
+      diffusion.DiffusionDistillationBatch.create(
+          student_batch={"target_ids": jnp.array([[0, 1]])},
+          teacher_logits=jnp.ones((1, 2, 2), dtype=jnp.float32),
+      )
+
+  def test_rejects_misaligned_teacher_logits(self):
+    with self.assertRaisesRegex(ValueError, "align with target_ids"):
+      diffusion.DiffusionDistillationBatch.create(
+          student_batch=_student_batch(),
+          teacher_logits=jnp.ones((1, 1, 2), dtype=jnp.float32),
+      )
+
+  def test_is_a_jax_pytree_with_aligned_teacher_logits(self):
+    teacher_logits = jnp.array([[[1.0, -1.0], [-0.5, 0.5]]], dtype=jnp.float32)
+    batch = diffusion.DiffusionDistillationBatch.create(
+        student_batch=_student_batch(),
+        teacher_logits=teacher_logits,
+    )
+
+    leaves, treedef = jax.tree.flatten(batch)
+    restored = jax.tree.unflatten(treedef, leaves)
+
+    self.assertLen(leaves, 4)
+    np.testing.assert_array_equal(restored.teacher_logits, teacher_logits)
+    np.testing.assert_array_equal(
+        restored.student_batch.target_ids, batch.student_batch.target_ids
+    )
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/distillation/diffusion_opd_test.py b/tests/distillation/diffusion_opd_test.py
new file mode 100644
index 000000000..2ac0b5ef2
--- /dev/null
+++ b/tests/distillation/diffusion_opd_test.py
@@ -0,0 +1,390 @@
+# 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.
+
+"""Tests for target-aligned diffusion on-policy distillation."""
+
+from absl.testing import absltest
+from flax import nnx
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+from tunix.diffusion import types as diffusion_types
+from tunix.distillation import diffusion as diffusion_distillation
+from tunix.distillation import diffusion_opd
+from tunix.sft import peft_trainer
+from tunix.sft import utils as sft_utils
+
+
+class _LogitModel(nnx.Module):
+
+  def __init__(self, logits):
+    self.logits = nnx.Param(logits)
+
+
+def _score_fn(
+    model: nnx.Module,
+    model_inputs: diffusion_types.ModelInputs,
+) -> jax.Array:
+  return model.logits[model_inputs["example_ids"]]
+
+
+def _batch(
+    *,
+    student_logits,
+    teacher_logits,
+    target_ids,
+    loss_weights,
+) -> tuple[_LogitModel, diffusion_distillation.DiffusionDistillationBatch]:
+  """Builds a test model and its aligned distillation batch."""
+
+  model = _LogitModel(jnp.asarray(student_logits))
+  student_batch = diffusion_types.DiffusionTokenBatch.create(
+      model_inputs={
+          "example_ids": jnp.arange(len(target_ids), dtype=jnp.int32)
+      },
+      target_ids=jnp.asarray(target_ids, dtype=jnp.int32),
+      loss_weights=jnp.asarray(loss_weights),
+  )
+  return model, diffusion_distillation.DiffusionDistillationBatch.create(
+      student_batch=student_batch,
+      teacher_logits=jnp.asarray(teacher_logits),
+  )
+
+
+def _loss_value(
+    model: nnx.Module,
+    batch: diffusion_distillation.DiffusionDistillationBatch,
+    **kwargs,
+) -> jax.Array:
+  return diffusion_opd.diffusion_opd_loss_fn(
+      model, batch, _score_fn, **kwargs
+  ).primary_loss.compute()
+
+
+class DiffusionOpdLossTest(absltest.TestCase):
+
+  def test_rejects_student_teacher_vocabulary_mismatch(self):
+    model, batch = _batch(
+        student_logits=jnp.ones((1, 2, 3), dtype=jnp.float32),
+        teacher_logits=jnp.ones((1, 2, 2), dtype=jnp.float32),
+        target_ids=[[0, 1]],
+        loss_weights=[[1.0, 1.0]],
+    )
+
+    with self.assertRaisesRegex(ValueError, "identical shapes"):
+      diffusion_opd.diffusion_opd_loss_fn(model, batch, _score_fn)
+
+  def test_uses_forward_kl_and_target_aligned_hard_loss(self):
+    student_logits = jnp.array([[[2.0, -1.0], [-0.5, 0.5]]], dtype=jnp.float32)
+    teacher_logits = jnp.array(
+        [[[-1.0, 2.0], [0.75, -0.75]]], dtype=jnp.float32
+    )
+    model, batch = _batch(
+        student_logits=student_logits,
+        teacher_logits=teacher_logits,
+        target_ids=[[1, 0]],
+        loss_weights=[[1.0, 0.5]],
+    )
+    temperature = 2.0
+
+    output = diffusion_opd.diffusion_opd_loss_fn(
+        model,
+        batch,
+        _score_fn,
+        temperature=temperature,
+        soft_loss_weight=0.75,
+        hard_loss_weight=0.25,
+    )
+    teacher_probs = jax.nn.softmax(teacher_logits / temperature, axis=-1)
+    expected_soft = optax.kl_divergence(
+        jax.nn.log_softmax(student_logits / temperature, axis=-1),
+        teacher_probs,
+    ) * (temperature**2)
+    expected_hard = optax.softmax_cross_entropy_with_integer_labels(
+        student_logits, batch.student_batch.target_ids
+    )
+    weights = batch.student_batch.loss_weights
+    soft_sum = jnp.sum(expected_soft * weights)
+    hard_sum = jnp.sum(expected_hard * weights)
+
+    self.assertIsInstance(output, sft_utils.LossOutput)
+    np.testing.assert_allclose(
+        output.primary_loss.unreduced_sum,
+        0.75 * soft_sum + 0.25 * hard_sum,
+        rtol=1e-6,
+    )
+    self.assertEqual(float(output.primary_loss.denominator), 1.5)
+    np.testing.assert_allclose(
+        output.aux_metrics["distill/soft_loss"].unreduced_sum,
+        soft_sum,
+        rtol=1e-6,
+    )
+    np.testing.assert_allclose(
+        output.aux_metrics["distill/hard_loss"].unreduced_sum,
+        hard_sum,
+        rtol=1e-6,
+    )
+
+  def test_teacher_logits_are_stop_gradient(self):
+    model, batch = _batch(
+        student_logits=jnp.array([[[2.0, -1.0]]], dtype=jnp.float32),
+        teacher_logits=jnp.array([[[-1.0, 2.0]]], dtype=jnp.float32),
+        target_ids=[[1]],
+        loss_weights=[[1.0]],
+    )
+
+    def loss_from_teacher(teacher_logits):
+      return _loss_value(
+          model,
+          batch.replace(teacher_logits=teacher_logits),  # pylint: disable=no-member
+      )
+
+    teacher_grads = jax.grad(loss_from_teacher)(batch.teacher_logits)
+    _, student_grads = nnx.value_and_grad(_loss_value)(model, batch)
+
+    np.testing.assert_array_equal(teacher_grads, jnp.zeros_like(teacher_grads))
+    self.assertTrue(
+        any(
+            bool(jnp.any(jnp.abs(leaf) > 0))
+            for leaf in jax.tree.leaves(student_grads)
+        )
+    )
+
+  def test_fractional_weights_apply_to_loss_and_metrics(self):
+    student_logits = jnp.array(
+        [[[1.0, 0.0], [0.0, 1.0], [2.0, -2.0]]], dtype=jnp.float32
+    )
+    teacher_logits = jnp.array(
+        [[[0.0, 1.0], [0.0, 1.0], [-2.0, 2.0]]], dtype=jnp.float32
+    )
+    model, batch = _batch(
+        student_logits=student_logits,
+        teacher_logits=teacher_logits,
+        target_ids=[[1, 1, 0]],
+        loss_weights=[[2.0, 0.25, 0.0]],
+    )
+
+    output = diffusion_opd.diffusion_opd_loss_fn(model, batch, _score_fn)
+    per_token = optax.kl_divergence(
+        jax.nn.log_softmax(student_logits, axis=-1),
+        jax.nn.softmax(teacher_logits, axis=-1),
+    )
+    expected_sum = jnp.sum(per_token * batch.student_batch.loss_weights)
+
+    np.testing.assert_allclose(
+        output.primary_loss.unreduced_sum, expected_sum, rtol=1e-6
+    )
+    self.assertEqual(float(output.primary_loss.denominator), 2.25)
+    self.assertEqual(
+        float(output.aux_metrics["distill/soft_loss"].denominator), 2.25
+    )
+
+  def test_zero_total_weight_has_zero_loss_and_gradients(self):
+    model, batch = _batch(
+        student_logits=jnp.array(
+            [[[jnp.inf, -jnp.inf], [jnp.nan, 0.0]]], dtype=jnp.float32
+        ),
+        teacher_logits=jnp.array(
+            [[[jnp.nan, 0.0], [jnp.inf, -jnp.inf]]], dtype=jnp.float32
+        ),
+        target_ids=[[0, 1]],
+        loss_weights=[[0.0, 0.0]],
+    )
+
+    @nnx.jit
+    def loss_and_grad(model, batch):
+      return nnx.value_and_grad(_loss_value)(model, batch)
+
+    output = diffusion_opd.diffusion_opd_loss_fn(
+        model,
+        batch,
+        _score_fn,
+        soft_loss_weight=0.5,
+        hard_loss_weight=0.5,
+    )
+    loss, grads = loss_and_grad(model, batch)
+
+    self.assertEqual(float(output.primary_loss.unreduced_sum), 0.0)
+    self.assertEqual(float(output.primary_loss.denominator), 0.0)
+    self.assertEqual(float(loss), 0.0)
+    for leaf in jax.tree.leaves(grads):
+      self.assertTrue(bool(jnp.all(jnp.isfinite(leaf))))
+      np.testing.assert_array_equal(leaf, jnp.zeros_like(leaf))
+
+  def test_inactive_invalid_logits_do_not_affect_active_loss_or_gradients(self):
+    model, batch = _batch(
+        student_logits=jnp.array(
+            [[[2.0, -1.0], [jnp.nan, jnp.inf]]], dtype=jnp.float32
+        ),
+        teacher_logits=jnp.array(
+            [[[-1.0, 2.0], [jnp.inf, jnp.nan]]], dtype=jnp.float32
+        ),
+        target_ids=[[1, -1]],
+        loss_weights=[[1.0, 0.0]],
+    )
+
+    loss, grads = nnx.value_and_grad(_loss_value)(model, batch)
+
+    self.assertTrue(bool(jnp.isfinite(loss)))
+    for leaf in jax.tree.leaves(grads):
+      self.assertTrue(bool(jnp.all(jnp.isfinite(leaf))))
+
+  def test_rejects_invalid_active_teacher_logits(self):
+    model, batch = _batch(
+        student_logits=jnp.array([[[1.0, -1.0]]], dtype=jnp.float32),
+        teacher_logits=jnp.array([[[jnp.nan, jnp.inf]]], dtype=jnp.float32),
+        target_ids=[[0]],
+        loss_weights=[[1.0]],
+    )
+
+    with self.assertRaisesRegex(ValueError, "active teacher_logits"):
+      diffusion_opd.diffusion_opd_loss_fn(model, batch, _score_fn)
+
+  def test_rejects_active_hard_targets_outside_vocabulary(self):
+    with self.assertRaisesRegex(ValueError, "smaller than vocabulary size"):
+      _batch(
+          student_logits=jnp.zeros((1, 1, 2), dtype=jnp.float32),
+          teacher_logits=jnp.zeros((1, 1, 2), dtype=jnp.float32),
+          target_ids=[[2]],
+          loss_weights=[[1.0]],
+      )
+
+  def test_disabled_soft_loss_ignores_invalid_teacher_values(self):
+    model, batch = _batch(
+        student_logits=jnp.array([[[1.0, -1.0]]], dtype=jnp.float32),
+        teacher_logits=jnp.array([[[jnp.nan, jnp.inf]]], dtype=jnp.float32),
+        target_ids=[[0]],
+        loss_weights=[[1.0]],
+    )
+
+    output = diffusion_opd.diffusion_opd_loss_fn(
+        model,
+        batch,
+        _score_fn,
+        soft_loss_weight=0.0,
+        hard_loss_weight=1.0,
+    )
+    expected = optax.softmax_cross_entropy_with_integer_labels(
+        model.logits[...], batch.student_batch.target_ids
+    )
+
+    self.assertTrue(bool(jnp.isfinite(output.primary_loss.compute())))
+    np.testing.assert_allclose(
+        output.primary_loss.unreduced_sum, jnp.sum(expected), rtol=1e-6
+    )
+    self.assertEqual(
+        float(output.aux_metrics["distill/soft_loss"].unreduced_sum), 0.0
+    )
+
+  def test_rejects_invalid_loss_configuration(self):
+    model, batch = _batch(
+        student_logits=jnp.zeros((1, 1, 2), dtype=jnp.float32),
+        teacher_logits=jnp.zeros((1, 1, 2), dtype=jnp.float32),
+        target_ids=[[0]],
+        loss_weights=[[1.0]],
+    )
+
+    with self.assertRaisesRegex(ValueError, "temperature must be positive"):
+      diffusion_opd.diffusion_opd_loss_fn(
+          model, batch, _score_fn, temperature=0.0
+      )
+    with self.assertRaisesRegex(ValueError, "must be nonnegative"):
+      diffusion_opd.diffusion_opd_loss_fn(
+          model, batch, _score_fn, soft_loss_weight=-1.0
+      )
+    with self.assertRaisesRegex(ValueError, "at least one"):
+      diffusion_opd.diffusion_opd_loss_fn(
+          model,
+          batch,
+          _score_fn,
+          soft_loss_weight=0.0,
+          hard_loss_weight=0.0,
+      )
+    for name, overrides in (
+        ("temperature", {"temperature": jnp.nan}),
+        ("temperature", {"temperature": jnp.inf}),
+        ("soft_loss_weight", {"soft_loss_weight": jnp.nan}),
+        ("soft_loss_weight", {"soft_loss_weight": jnp.inf}),
+        ("hard_loss_weight", {"hard_loss_weight": jnp.nan}),
+        ("hard_loss_weight", {"hard_loss_weight": jnp.inf}),
+    ):
+      with self.subTest(name=name, overrides=overrides):
+        with self.assertRaisesRegex(ValueError, "finite real scalar"):
+          diffusion_opd.diffusion_opd_loss_fn(
+              model,
+              batch,
+              _score_fn,
+              **overrides,
+          )
+
+  def test_configure_prepared_diffusion_opd_wires_public_trainer(self):
+    model = _LogitModel(
+        jnp.array([[[1.0, -1.0], [-0.5, 0.5]]], dtype=jnp.float32)
+    )
+    trainer = peft_trainer.PeftTrainer(
+        model,
+        optax.sgd(0.1),
+        peft_trainer.TrainingConfig(eval_every_n_steps=100, max_steps=1),
+    )
+    raw_batch = {
+        "target_ids": [[1, 0]],
+        "loss_weights": [[1.0, 0.5]],
+        "teacher_logits": [[[-1.0, 1.0], [1.0, -1.0]]],
+    }
+    adapter_inputs = []
+
+    def batch_adapter(value):
+      adapter_inputs.append(value)
+      student_batch = diffusion_types.DiffusionTokenBatch.create(
+          model_inputs={"example_ids": jnp.array([0], dtype=jnp.int32)},
+          target_ids=jnp.asarray(value["target_ids"], dtype=jnp.int32),
+          loss_weights=jnp.asarray(value["loss_weights"]),
+      )
+      return diffusion_distillation.DiffusionDistillationBatch.create(
+          student_batch=student_batch,
+          teacher_logits=jnp.asarray(value["teacher_logits"]),
+      )
+
+    configured_trainer = diffusion_opd.configure_prepared_diffusion_opd(
+        trainer,
+        batch_adapter,
+        _score_fn,
+        temperature=2.0,
+        soft_loss_weight=0.8,
+        hard_loss_weight=0.2,
+    )
+    original_logits = jnp.copy(model.logits[...])
+    loss, aux, grad_norm = configured_trainer.create_train_step_fn()(
+        model,
+        configured_trainer.optimizer,
+        configured_trainer.grad_accumulator,
+        raw_batch,
+        jnp.asarray(True),
+    )
+
+    self.assertIs(configured_trainer, trainer)
+    self.assertLen(adapter_inputs, 1)
+    self.assertIs(adapter_inputs[0], raw_batch)
+    self.assertTrue(bool(jnp.isfinite(loss)))
+    self.assertIsInstance(aux, sft_utils.LossOutput)
+    self.assertCountEqual(
+        aux.aux_metrics, ["distill/soft_loss", "distill/hard_loss"]
+    )
+    self.assertGreater(float(grad_norm), 0.0)
+    self.assertFalse(bool(jnp.array_equal(model.logits[...], original_logits)))
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/experimental/train/peft_trainer_v2_test.py b/tests/experimental/train/peft_trainer_v2_test.py
index 28fa0baf6..047bded97 100644
--- a/tests/experimental/train/peft_trainer_v2_test.py
+++ b/tests/experimental/train/peft_trainer_v2_test.py
@@ -20,6 +20,7 @@
 import tempfile
 from typing import Any, Tuple
 from unittest import mock
+
 from absl.testing import absltest
 from absl.testing import parameterized
 import chex
@@ -33,8 +34,8 @@
 from tunix.experimental.train import peft_trainer_v2
 from tunix.sft import checkpoint_manager
 from tunix.sft import hooks
-from tunix.sft import profiler
 from tunix.sft import peft_trainer
+from tunix.sft import profiler
 from tunix.sft import utils as sft_utils
 from tunix.tests import test_common as tc
 from tunix.utils import compat
@@ -221,7 +222,9 @@ def create_model_and_optimizer():
     # the model and optimizer states are the same as the trained ones.
     new_model, new_optimizer = create_model_and_optimizer()
 
-    resumed_trainer = peft_trainer_v2.PeftTrainer(new_model, new_optimizer, config)
+    resumed_trainer = peft_trainer_v2.PeftTrainer(
+        new_model, new_optimizer, config
+    )
     resumed_model_state = nnx.state(
         resumed_trainer.model, nnx.LoRAParam if enable_lora else nnx.Param
     )
@@ -344,7 +347,9 @@ def test_dist_training(self):
     # compare with unsharded model
     rngs = nnx.Rngs(0)
     unsharded_model = tc.ToyTransformer(config=tc.ModelConfig(), rngs=rngs)
-    trainer = peft_trainer_v2.PeftTrainer(unsharded_model, optax.sgd(1e-3), config)
+    trainer = peft_trainer_v2.PeftTrainer(
+        unsharded_model, optax.sgd(1e-3), config
+    )
     trainer = trainer.with_gen_model_input_fn(dummy_gen_model_input_fn)
     trainer.train(self.train_ds, self.eval_ds)
     unsharded_variables = nnx.state(unsharded_model, nnx.Param)
@@ -531,6 +536,10 @@ def test_checkpointing(
     )
     trainer = peft_trainer_v2.PeftTrainer(model, optax.sgd(1e-3), config)
     trainer = trainer.with_gen_model_input_fn(dummy_gen_model_input_fn)
+    checkpoint_metadata = {'training_contract': 'test-v1'}
+    trainer.custom_checkpoint_metadata = mock.Mock(
+        return_value=checkpoint_metadata
+    )
 
     train_ds = eval_ds = dummy_datasets(batch_size=2, repeat=1)  # 4 batches
     trainer.train(train_ds, eval_ds)
@@ -551,7 +560,7 @@ def test_checkpointing(
                     mock.ANY,
                     mock.ANY,
                     save_only_lora_params=True,
-                    custom_metadata={},
+                    custom_metadata=checkpoint_metadata,
                 )
                 for i in expected_save_steps
             ],
@@ -561,6 +570,7 @@ def test_checkpointing(
                 mock.ANY,
                 mock.ANY,
                 save_only_lora_params=True,
+                custom_metadata=checkpoint_metadata,
                 force=True,
             ),
             mock.call.close(),
@@ -622,9 +632,7 @@ def test_explicit_save_checkpoint(self, mock_checkpoint_manager_init):
     trainer = peft_trainer_v2.PeftTrainer(model, optax.sgd(1e-3), config)
 
     custom_metadata = {'global_step': 42, 'role': 'actor'}
-    trainer.save_checkpoint(
-        metadata=custom_metadata, step=10, force=True
-    )
+    trainer.save_checkpoint(metadata=custom_metadata, step=10, force=True)
 
     mock_cm.save.assert_called_once_with(
         10,
@@ -784,7 +792,9 @@ def test_parity_with_v1_trainer(self):
     optimizer_v2 = optax.inject_hyperparams(optax.sgd)(
         learning_rate=optax.constant_schedule(TEST_LEARNING_RATE)
     )
-    config_v2 = peft_trainer_v2.TrainingConfig(eval_every_n_steps=2, max_steps=1)
+    config_v2 = peft_trainer_v2.TrainingConfig(
+        eval_every_n_steps=2, max_steps=1
+    )
     trainer_v2 = peft_trainer_v2.PeftTrainer(model_v2, optimizer_v2, config_v2)
     trainer_v2 = trainer_v2.with_gen_model_input_fn(dummy_gen_model_input_fn)
     trainer_v2.train(train_ds, None)
diff --git a/tests/rl/diffusion_test.py b/tests/rl/diffusion_test.py
new file mode 100644
index 000000000..6f35f3404
--- /dev/null
+++ b/tests/rl/diffusion_test.py
@@ -0,0 +1,241 @@
+# 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.
+
+"""Tests for prepared diffusion policy scoring."""
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from flax import nnx
+import jax
+import jax.numpy as jnp
+import numpy as np
+from tunix.diffusion import types as diffusion_types
+from tunix.rl import diffusion
+from tunix.rl import rl_cluster
+from tunix.rl.rollout import base_rollout
+
+
+class _LogitModel(nnx.Module):
+
+  def __init__(self, logits):
+    self.logits = nnx.Param(logits)
+
+
+def _score_fn(model, model_inputs):
+  return model.logits[model_inputs["example_ids"]]
+
+
+def _batch(*, target_ids, loss_weights):
+  batch_size = len(target_ids)
+  return diffusion_types.DiffusionTokenBatch.create(
+      model_inputs={
+          "example_ids": jnp.arange(batch_size, dtype=jnp.int32),
+      },
+      target_ids=jnp.asarray(target_ids, dtype=jnp.int32),
+      loss_weights=jnp.asarray(loss_weights, dtype=jnp.float32),
+  )
+
+
+def _rollout_output(diffusion_batch=None):
+  batch_size = (
+      1 if diffusion_batch is None else diffusion_batch.target_ids.shape[0]
+  )
+  return base_rollout.RolloutOutput(
+      text=["generated"] * batch_size,
+      logits=None,
+      tokens=[np.zeros((1,), dtype=np.int32)] * batch_size,
+      left_padded_prompt_tokens=np.zeros((batch_size, 1), dtype=np.int32),
+      logprobs=None,
+      diffusion_batch=diffusion_batch,
+  )
+
+
+class DiffusionPolicyScoringTest(parameterized.TestCase):
+
+  def test_returns_temperature_scaled_completion_aligned_scores_and_entropy(
+      self,
+  ):
+    logits = jnp.array(
+        [
+            [[2.0, 0.0, -1.0], [0.0, 1.0, 3.0]],
+            [[-2.0, 1.0, 0.0], [4.0, 0.0, -1.0]],
+        ],
+        dtype=jnp.float32,
+    )
+    batch = _batch(
+        target_ids=[[0, 2], [1, 0]],
+        loss_weights=[[1.0, 1.0], [1.0, 1.0]],
+    )
+
+    logps, entropy = diffusion.compute_diffusion_per_token_logps(
+        _LogitModel(logits),
+        batch,
+        _score_fn,
+        temperature=2.0,
+        stop_gradient=False,
+        return_entropy=True,
+    )
+    expected_log_probs = jax.nn.log_softmax(logits / 2.0, axis=-1)
+    expected_logps = jnp.take_along_axis(
+        expected_log_probs, batch.target_ids[..., None], axis=-1
+    )[..., 0]
+    expected_entropy = -jnp.sum(
+        jnp.exp(expected_log_probs) * expected_log_probs, axis=-1
+    )
+
+    self.assertEqual(logps.shape, batch.target_ids.shape)
+    self.assertEqual(entropy.shape, batch.target_ids.shape)
+    np.testing.assert_allclose(logps, expected_logps, rtol=1e-6)
+    np.testing.assert_allclose(entropy, expected_entropy, rtol=1e-6)
+
+  def test_inactive_positions_sanitize_sentinel_targets_and_nonfinite_logits(
+      self,
+  ):
+    logits = jnp.array(
+        [[[2.0, -1.0], [jnp.nan, jnp.inf], [jnp.inf, -jnp.inf]]],
+        dtype=jnp.float32,
+    )
+    batch = _batch(
+        target_ids=[[0, -1, 99]],
+        loss_weights=[[1.0, 0.0, 0.0]],
+    )
+
+    logps, entropy = diffusion.compute_diffusion_per_token_logps(
+        _LogitModel(logits),
+        batch,
+        _score_fn,
+        return_entropy=True,
+    )
+
+    self.assertTrue(bool(jnp.all(jnp.isfinite(logps))))
+    self.assertTrue(bool(jnp.all(jnp.isfinite(entropy))))
+    np.testing.assert_array_equal(logps[0, 1:], jnp.zeros((2,)))
+    np.testing.assert_array_equal(entropy[0, 1:], jnp.zeros((2,)))
+
+  def test_entropy_remains_finite_when_active_logits_mask_a_token(self):
+    logits = jnp.array([[[2.0, 0.0, -jnp.inf]]], dtype=jnp.float32)
+    batch = _batch(target_ids=[[0]], loss_weights=[[1.0]])
+
+    logps, entropy = diffusion.compute_diffusion_per_token_logps(
+        _LogitModel(logits),
+        batch,
+        _score_fn,
+        return_entropy=True,
+    )
+
+    self.assertTrue(bool(jnp.all(jnp.isfinite(logps))))
+    self.assertTrue(bool(jnp.all(jnp.isfinite(entropy))))
+
+  def test_stop_gradient_controls_policy_gradient(self):
+    batch = _batch(
+        target_ids=[[0, 1]],
+        loss_weights=[[1.0, 1.0]],
+    )
+
+    def score_sum(model, stop_gradient):
+      return jnp.sum(
+          diffusion.compute_diffusion_per_token_logps(
+              model,
+              batch,
+              _score_fn,
+              stop_gradient=stop_gradient,
+          )
+      )
+
+    trainable_model = _LogitModel(
+        jnp.array([[[1.0, -1.0], [-1.0, 1.0]]], dtype=jnp.float32)
+    )
+    detached_model = _LogitModel(jnp.copy(trainable_model.logits[...]))
+    trainable_grads = nnx.grad(lambda model: score_sum(model, False))(
+        trainable_model
+    )
+    detached_grads = nnx.grad(lambda model: score_sum(model, True))(
+        detached_model
+    )
+
+    self.assertGreater(float(jnp.linalg.norm(trainable_grads.logits[...])), 0)
+    np.testing.assert_array_equal(
+        detached_grads.logits[...], jnp.zeros_like(detached_grads.logits[...])
+    )
+
+  @parameterized.parameters(0.0, -1.0, float("inf"), float("nan"), True, "1")
+  def test_rejects_invalid_temperature(self, temperature):
+    model = _LogitModel(jnp.zeros((1, 1, 2), dtype=jnp.float32))
+    batch = _batch(target_ids=[[0]], loss_weights=[[1.0]])
+
+    with self.assertRaisesRegex(ValueError, "temperature"):
+      diffusion.compute_diffusion_per_token_logps(
+          model, batch, _score_fn, temperature=temperature
+      )
+
+
+class DiffusionRolloutMergeTest(absltest.TestCase):
+
+  def test_merges_prepared_batches_on_leading_dimension(self):
+    first = _batch(target_ids=[[1, 2]], loss_weights=[[1.0, 1.0]])
+    second = _batch(target_ids=[[3, 4]], loss_weights=[[1.0, 0.0]])
+
+    merged = rl_cluster._merge_diffusion_batches(  # pylint: disable=protected-access
+        [_rollout_output(first), _rollout_output(second)]
+    )
+
+    self.assertIsNotNone(merged)
+    np.testing.assert_array_equal(merged.target_ids, [[1, 2], [3, 4]])
+    np.testing.assert_array_equal(merged.loss_weights, [[1.0, 1.0], [1.0, 0.0]])
+    np.testing.assert_array_equal(merged.model_inputs["example_ids"], [0, 0])
+
+  def test_allows_ar_microbatches_to_omit_diffusion_metadata(self):
+    self.assertIsNone(
+        rl_cluster._merge_diffusion_batches(  # pylint: disable=protected-access
+            [_rollout_output(), _rollout_output()]
+        )
+    )
+
+  def test_rejects_mixed_diffusion_metadata(self):
+    batch = _batch(target_ids=[[1]], loss_weights=[[1.0]])
+
+    with self.assertRaisesRegex(ValueError, "either all provide"):
+      rl_cluster._merge_diffusion_batches(  # pylint: disable=protected-access
+          [_rollout_output(batch), _rollout_output()]
+      )
+
+  def test_rejects_different_model_input_pytrees(self):
+    first = _batch(target_ids=[[1]], loss_weights=[[1.0]])
+    second = diffusion_types.DiffusionTokenBatch.create(
+        model_inputs={"other_ids": jnp.array([0], dtype=jnp.int32)},
+        target_ids=jnp.array([[2]], dtype=jnp.int32),
+        loss_weights=jnp.array([[1.0]], dtype=jnp.float32),
+    )
+
+    with self.assertRaisesRegex(ValueError, "same pytree structure"):
+      rl_cluster._merge_diffusion_batches(  # pylint: disable=protected-access
+          [_rollout_output(first), _rollout_output(second)]
+      )
+
+  def test_rejects_batch_size_that_does_not_match_rollout(self):
+    batch = _batch(
+        target_ids=[[1], [2]],
+        loss_weights=[[1.0], [1.0]],
+    )
+    output = _rollout_output(batch)
+    output.text = ["only one generated example"]
+
+    with self.assertRaisesRegex(ValueError, "batch size must match"):
+      rl_cluster._merge_diffusion_batches(  # pylint: disable=protected-access
+          [output]
+      )
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/sft/diffusion_test.py b/tests/sft/diffusion_test.py
new file mode 100644
index 000000000..7f33e33c7
--- /dev/null
+++ b/tests/sft/diffusion_test.py
@@ -0,0 +1,279 @@
+# 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.
+
+"""Tests for diffusion supervised fine-tuning loss."""
+
+from absl.testing import absltest
+from flax import nnx
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+from tunix.diffusion import types as diffusion_types
+from tunix.sft import diffusion
+from tunix.sft import peft_trainer
+from tunix.sft import utils as sft_utils
+
+
+class _LogitModel(nnx.Module):
+
+  def __init__(self, logits):
+    self.logits = nnx.Param(logits)
+
+
+def _score_fn(
+    model: nnx.Module,
+    model_inputs: diffusion_types.ModelInputs,
+) -> jax.Array:
+  return model.logits[model_inputs["example_ids"]]
+
+
+def _batch(
+    *,
+    example_ids,
+    target_ids,
+    loss_weights,
+) -> diffusion_types.DiffusionTokenBatch:
+  return diffusion_types.DiffusionTokenBatch.create(
+      model_inputs={"example_ids": jnp.asarray(example_ids, dtype=jnp.int32)},
+      target_ids=jnp.asarray(target_ids, dtype=jnp.int32),
+      loss_weights=jnp.asarray(loss_weights),
+  )
+
+
+def _loss_value(
+    model: nnx.Module,
+    batch: diffusion_types.DiffusionTokenBatch,
+) -> jax.Array:
+  return diffusion.diffusion_loss_fn(
+      model, batch, _score_fn
+  ).primary_loss.compute()
+
+
+class DiffusionLossTest(absltest.TestCase):
+
+  def test_configure_diffusion_sft_wires_trainer(self):
+    model = _LogitModel(
+        jnp.array([[[1.0, -1.0], [-0.5, 0.5]]], dtype=jnp.float32)
+    )
+    trainer = peft_trainer.PeftTrainer(
+        model,
+        optax.sgd(0.1),
+        peft_trainer.TrainingConfig(eval_every_n_steps=100, max_steps=1),
+    )
+    raw_batch = {
+        "example_ids": [0],
+        "target_ids": [[1, 0]],
+        "loss_weights": [[1.0, 1.0]],
+    }
+    adapter_inputs = []
+
+    def batch_adapter(value):
+      adapter_inputs.append(value)
+      return _batch(
+          example_ids=value["example_ids"],
+          target_ids=value["target_ids"],
+          loss_weights=value["loss_weights"],
+      )
+
+    configured_trainer = diffusion.configure_diffusion_sft(
+        trainer, batch_adapter, _score_fn
+    )
+    original_logits = jnp.copy(model.logits[...])
+    loss, aux, grad_norm = configured_trainer.create_train_step_fn()(
+        model,
+        configured_trainer.optimizer,
+        configured_trainer.grad_accumulator,
+        raw_batch,
+        jnp.asarray(True),
+    )
+
+    self.assertIs(configured_trainer, trainer)
+    self.assertLen(adapter_inputs, 1)
+    self.assertIs(adapter_inputs[0], raw_batch)
+    self.assertTrue(bool(jnp.isfinite(loss)))
+    self.assertIsInstance(aux, sft_utils.LossOutput)
+    self.assertEmpty(aux.aux_metrics)
+    self.assertGreater(float(grad_norm), 0.0)
+    self.assertFalse(bool(jnp.array_equal(model.logits[...], original_logits)))
+
+  def test_uses_target_aligned_logits_without_shifting(self):
+    logits = jnp.array(
+        [[[8.0, 0.0, 0.0], [0.0, 8.0, 0.0], [0.0, 0.0, 8.0]]],
+        dtype=jnp.float32,
+    )
+    batch = _batch(
+        example_ids=[0],
+        target_ids=[[0, 1, 2]],
+        loss_weights=[[1.0, 0.5, 0.0]],
+    )
+
+    output = diffusion.diffusion_loss_fn(_LogitModel(logits), batch, _score_fn)
+    expected_token_losses = optax.softmax_cross_entropy_with_integer_labels(
+        logits, batch.target_ids
+    )
+    expected_sum = jnp.sum(
+        expected_token_losses * batch.loss_weights, dtype=jnp.float32
+    )
+
+    self.assertIsInstance(output, sft_utils.LossOutput)
+    np.testing.assert_allclose(
+        output.primary_loss.unreduced_sum, expected_sum, rtol=1e-6
+    )
+    self.assertEqual(float(output.primary_loss.denominator), 1.5)
+    self.assertEmpty(output.aux_metrics)
+
+  def test_cross_entropy_and_metrics_are_float32(self):
+    logits = jnp.array([[[2.0, -1.0], [-2.0, 3.0]]], dtype=jnp.bfloat16)
+    batch = _batch(
+        example_ids=[0],
+        target_ids=[[0, 1]],
+        loss_weights=jnp.array([[1, 2]], dtype=jnp.int32),
+    )
+
+    metric = diffusion.diffusion_loss_fn(
+        _LogitModel(logits), batch, _score_fn
+    ).primary_loss
+
+    self.assertEqual(metric.unreduced_sum.dtype, jnp.float32)
+    self.assertEqual(metric.denominator.dtype, jnp.float32)
+    self.assertEqual(metric.compute().dtype, jnp.float32)
+
+  def test_zero_total_weight_has_zero_loss_and_gradient(self):
+    model = _LogitModel(
+        jnp.array([[[jnp.inf, -jnp.inf], [jnp.nan, 0.5]]], dtype=jnp.float32)
+    )
+    batch = _batch(
+        example_ids=[0],
+        target_ids=[[0, 1]],
+        loss_weights=[[0.0, 0.0]],
+    )
+
+    @nnx.jit
+    def loss_and_grad(model, batch):
+      return nnx.value_and_grad(_loss_value)(model, batch)
+
+    metric = diffusion.diffusion_loss_fn(model, batch, _score_fn).primary_loss
+    loss, grads = loss_and_grad(model, batch)
+
+    self.assertEqual(float(metric.unreduced_sum), 0.0)
+    self.assertEqual(float(metric.denominator), 0.0)
+    self.assertEqual(float(loss), 0.0)
+    for leaf in jax.tree.leaves(grads):
+      self.assertTrue(bool(jnp.all(jnp.isfinite(leaf))))
+      np.testing.assert_array_equal(leaf, jnp.zeros_like(leaf))
+
+  def test_jitted_gradients_are_finite_and_nonzero(self):
+    model = _LogitModel(
+        jnp.array([[[1.0, -1.0], [-0.5, 0.5]]], dtype=jnp.float32)
+    )
+    batch = _batch(
+        example_ids=[0],
+        target_ids=[[1, 0]],
+        loss_weights=[[1.0, 1.0]],
+    )
+
+    @nnx.jit
+    def loss_and_grad(model, batch):
+      return nnx.value_and_grad(_loss_value)(model, batch)
+
+    loss, grads = loss_and_grad(model, batch)
+    gradient_leaves = jax.tree.leaves(grads)
+
+    self.assertTrue(bool(jnp.isfinite(loss)))
+    self.assertTrue(
+        any(bool(jnp.any(jnp.abs(leaf) > 0)) for leaf in gradient_leaves)
+    )
+    for leaf in gradient_leaves:
+      self.assertTrue(bool(jnp.all(jnp.isfinite(leaf))))
+
+  def test_equal_weight_microbatches_match_full_batch(self):
+    logits = jnp.array(
+        [
+            [[2.0, -1.0], [0.5, -0.5]],
+            [[-1.0, 2.0], [-0.5, 0.5]],
+            [[0.0, 1.0], [1.5, -0.5]],
+            [[1.0, 0.0], [-1.5, 0.5]],
+        ],
+        dtype=jnp.float32,
+    )
+    targets = jnp.array([[0, 1], [1, 0], [0, 0], [1, 1]])
+    weights = jnp.ones((4, 2), dtype=jnp.float32)
+    full_batch = _batch(
+        example_ids=jnp.arange(4),
+        target_ids=targets,
+        loss_weights=weights,
+    )
+    microbatches = [
+        _batch(
+            example_ids=jnp.arange(start, start + 2),
+            target_ids=targets[start : start + 2],
+            loss_weights=weights[start : start + 2],
+        )
+        for start in (0, 2)
+    ]
+    model = _LogitModel(logits)
+
+    full_output = diffusion.diffusion_loss_fn(model, full_batch, _score_fn)
+    micro_outputs = [
+        diffusion.diffusion_loss_fn(model, batch, _score_fn)
+        for batch in microbatches
+    ]
+    _, full_grads = nnx.value_and_grad(_loss_value)(model, full_batch)
+    micro_grads = [
+        nnx.value_and_grad(_loss_value)(model, batch)[1]
+        for batch in microbatches
+    ]
+    mean_micro_grads = jax.tree.map(
+        lambda first, second: (first + second) / 2,
+        micro_grads[0],
+        micro_grads[1],
+    )
+
+    np.testing.assert_allclose(
+        full_output.primary_loss.unreduced_sum,
+        sum(output.primary_loss.unreduced_sum for output in micro_outputs),
+        rtol=1e-6,
+    )
+    self.assertEqual(
+        float(full_output.primary_loss.denominator),
+        sum(float(output.primary_loss.denominator) for output in micro_outputs),
+    )
+    jax.tree.map(
+        lambda actual, expected: np.testing.assert_allclose(
+            actual, expected, rtol=1e-6, atol=1e-6
+        ),
+        mean_micro_grads,
+        full_grads,
+    )
+
+  def test_rejects_misaligned_logits(self):
+    batch = _batch(
+        example_ids=[0],
+        target_ids=[[0, 1]],
+        loss_weights=[[1.0, 1.0]],
+    )
+
+    def misaligned_score_fn(model, model_inputs):
+      del model, model_inputs
+      return jnp.ones((1, 1, 2), dtype=jnp.float32)
+
+    with self.assertRaisesRegex(ValueError, "align with target_ids"):
+      diffusion.diffusion_loss_fn(
+          _LogitModel(jnp.zeros((1, 2, 2))), batch, misaligned_score_fn
+      )
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/sft/peft_trainer_test.py b/tests/sft/peft_trainer_test.py
index 6e4afddf4..45479ebc0 100644
--- a/tests/sft/peft_trainer_test.py
+++ b/tests/sft/peft_trainer_test.py
@@ -20,6 +20,7 @@
 import tempfile
 from typing import Any, Tuple
 from unittest import mock
+
 from absl.testing import absltest
 from absl.testing import parameterized
 import chex
@@ -47,6 +48,7 @@
 # Set Precision to highest for numeric stability across different hardware.
 jax.config.update('jax_default_matmul_precision', 'highest')
 
+
 def create_sharded_model(model_ctor, rngs, mesh):
   @nnx.jit(static_argnums=(0,))
   def _create_sharded_model(model_ctor, rngs):
@@ -569,6 +571,10 @@ def test_checkpointing(
     )
     trainer = peft_trainer.PeftTrainer(model, optax.sgd(1e-3), config)
     trainer = trainer.with_gen_model_input_fn(dummy_gen_model_input_fn)
+    checkpoint_metadata = {'training_contract': 'test-v1'}
+    trainer.custom_checkpoint_metadata = mock.Mock(
+        return_value=checkpoint_metadata
+    )
 
     train_ds = eval_ds = dummy_datasets(batch_size=2, repeat=1)  # 4 batches
     trainer.train(train_ds, eval_ds)
@@ -589,7 +595,7 @@ def test_checkpointing(
                     mock.ANY,
                     mock.ANY,
                     save_only_lora_params=True,
-                    custom_metadata={},
+                    custom_metadata=checkpoint_metadata,
                 )
                 for i in expected_save_steps
             ],
@@ -599,6 +605,7 @@ def test_checkpointing(
                 mock.ANY,
                 mock.ANY,
                 save_only_lora_params=True,
+                custom_metadata=checkpoint_metadata,
                 force=True,
             ),
             mock.call.close(),
@@ -702,10 +709,107 @@ def _post_process_eval_step(self, aux):
     # Since eval_ds is length 2, it evaluates at step 2.
     self.assertEqual(eval_invoke, {'foo': 8.0, 'bar': 12.0})
 
+  def test_loss_output_metrics_support_plain_scalars(self):
+    def custom_loss_fn(
+        model: nnx.Module,
+        input_tokens: jax.Array,
+        input_mask: jax.Array,
+        positions: jax.Array,
+        attention_mask: jax.Array,
+        images: jax.Array | None = None,
+    ) -> utils.LossOutput:
+      del model, input_tokens, input_mask, positions, attention_mask, images
+      return utils.LossOutput(
+          primary_loss=utils.WeightedMetric(
+              jnp.array(2.0, dtype=jnp.float32),
+              jnp.array(2.0, dtype=jnp.float32),
+          ),
+          aux_metrics={
+              'weighted': utils.WeightedMetric(
+                  jnp.array(10.0, dtype=jnp.float32),
+                  jnp.array(5.0, dtype=jnp.float32),
+              ),
+              'plain': jnp.array(4.0, dtype=jnp.float32),
+          },
+      )
+
+    config = peft_trainer.TrainingConfig(eval_every_n_steps=2, max_steps=100)
+    model = tc.ToyTransformer(config=tc.ModelConfig(), rngs=nnx.Rngs(0))
+    trainer = peft_trainer.PeftTrainer(model, optax.sgd(1e-3), config)
+    trainer = trainer.with_gen_model_input_fn(
+        dummy_gen_model_input_fn
+    ).with_loss_fn(custom_loss_fn)
+
+    trainer.train(self.train_ds, self.eval_ds)
+
+    self.assertEqual(
+        trainer.metrics_logger.get_metric('', 'weighted', 'train'), 2.0
+    )
+    self.assertEqual(
+        trainer.metrics_logger.get_metric('', 'plain', 'train'), 4.0
+    )
+    self.assertEqual(
+        trainer.metrics_logger.get_metric('', 'weighted', 'eval'), 2.0
+    )
+    self.assertEqual(
+        trainer.metrics_logger.get_metric('', 'plain', 'eval'), 4.0
+    )
+
+  def test_metrics_buffer_rejects_mixed_metric_kinds(self):
+    weighted = utils.WeightedMetric(jnp.array(1.0), jnp.array(1.0))
+    buffer = peft_trainer.MetricsBuffer(
+        step=0,
+        losses=[jnp.array(1.0), weighted],
+    )
+    with self.assertRaisesRegex(TypeError, 'must not mix'):
+      _ = buffer.loss
+
+    model = tc.ToyTransformer(config=tc.ModelConfig(), rngs=nnx.Rngs(0))
+    trainer = peft_trainer.PeftTrainer(
+        model,
+        optax.sgd(1e-3),
+        peft_trainer.TrainingConfig(eval_every_n_steps=100, max_steps=100),
+    )
+    buffer = trainer._buffer_metrics(
+        None,
+        loss=jnp.array(1.0),
+        step=0,
+        additional_metrics={'metric': (jnp.array(1.0), np.mean)},
+    )
+    with self.assertRaisesRegex(TypeError, "additional metric 'metric'"):
+      trainer._buffer_metrics(
+          buffer,
+          loss=jnp.array(1.0),
+          step=0,
+          additional_metrics={
+              'metric': (weighted, peft_trainer._weighted_metric_mean),
+          },
+      )
+
+  def test_weighted_metric_mean_preserves_denominator_bounds(self):
+    metrics = [
+        utils.WeightedMetric(
+            jnp.array(3.0), jnp.array(0.0), eps=1.0, min_denom=2.0
+        ),
+        utils.WeightedMetric(
+            jnp.array(1.0), jnp.array(0.0), eps=1.0, min_denom=2.0
+        ),
+    ]
+    self.assertEqual(peft_trainer._weighted_metric_mean(metrics), 2.0)
+
+    inconsistent = [
+        metrics[0],
+        utils.WeightedMetric(
+            jnp.array(1.0), jnp.array(0.0), eps=1.0, min_denom=3.0
+        ),
+    ]
+    with self.assertRaisesRegex(ValueError, 'consistent denominator bounds'):
+      peft_trainer._weighted_metric_mean(inconsistent)
+
   def test_loss_output_gradient_scaling(self):
-    # Covers the manual gradient scaling in _train_step: grad(unreduced_sum) *
-    # (1/d) must equal grad(sum/d). Uses a parameter-dependent loss because the
-    # original test's constant loss has zero gradient and can't exercise it.
+    # Covers denominator-aware accumulation in _train_step: grad(sum) / d must
+    # equal grad(sum / d). Uses a parameter-dependent loss because the original
+    # test's constant loss has zero gradient and can't exercise it.
     def param_dependent_sum(model, input_tokens, positions):
       logits, _ = model(input_tokens, positions)
       return jnp.sum(logits)  # depends on the parameters
@@ -724,7 +828,11 @@ def unreduced_loss_fn(
 
     def make_reduced_loss_fn(denominator):
       def reduced_loss_fn(
-          model, input_tokens, input_mask, positions, attention_mask,
+          model,
+          input_tokens,
+          input_mask,
+          positions,
+          attention_mask,
           images=None,
       ):
         del input_mask, attention_mask, images
@@ -770,6 +878,145 @@ def max_abs_diff(a, b):
     # the scale is applied.
     self.assertGreater(max_abs_diff(params_a, params_b1), 1e-6)
 
+  def test_loss_output_accumulates_variable_denominators_globally(self):
+    inputs = jnp.array(
+        [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [2.0, -1.0]],
+        dtype=jnp.float32,
+    )
+    targets = jnp.array([[0.5], [-0.5], [1.0], [-1.0]], dtype=jnp.float32)
+    weights = jnp.array([0.0, 0.0, 0.2, 0.3], dtype=jnp.float32)
+
+    def weighted_loss_sum(model, x, y, example_weights):
+      per_example = jnp.sum(jnp.square(model(x) - y), axis=-1)
+      return jnp.sum(per_example * example_weights)
+
+    def unreduced_loss_fn(model, x, y, example_weights):
+      return utils.LossOutput(
+          primary_loss=utils.WeightedMetric(
+              weighted_loss_sum(model, x, y, example_weights),
+              jnp.sum(example_weights),
+          ),
+          aux_metrics={},
+      )
+
+    def reduced_loss_fn(model, x, y, example_weights):
+      return weighted_loss_sum(model, x, y, example_weights) / jnp.sum(
+          example_weights
+      )
+
+    config = peft_trainer.TrainingConfig(
+        eval_every_n_steps=100,
+        max_steps=1,
+        gradient_accumulation_steps=2,
+    )
+
+    accumulated_model = nnx.Linear(2, 1, rngs=nnx.Rngs(0))
+    accumulated_trainer = peft_trainer.PeftTrainer(
+        accumulated_model, optax.sgd(0.1), config
+    ).with_loss_fn(unreduced_loss_fn)
+    accumulated_grads = peft_trainer.GradientAccumulator(
+        accumulated_model, nnx.Param
+    )
+    accumulated_train_step = accumulated_trainer.create_train_step_fn()
+
+    for index, is_update_step in enumerate((False, True)):
+      batch_slice = slice(index * 2, (index + 1) * 2)
+      accumulated_train_step(
+          accumulated_model,
+          accumulated_trainer.optimizer,
+          accumulated_grads,
+          {
+              'x': inputs[batch_slice],
+              'y': targets[batch_slice],
+              'example_weights': weights[batch_slice],
+          },
+          jnp.asarray(is_update_step),
+      )
+
+    reference_model = nnx.Linear(2, 1, rngs=nnx.Rngs(0))
+    reference_trainer = peft_trainer.PeftTrainer(
+        reference_model, optax.sgd(0.1), config
+    ).with_loss_fn(reduced_loss_fn)
+    reference_grads = peft_trainer.GradientAccumulator(
+        reference_model, nnx.Param
+    )
+    reference_trainer.create_train_step_fn()(
+        reference_model,
+        reference_trainer.optimizer,
+        reference_grads,
+        {
+            'x': inputs,
+            'y': targets,
+            'example_weights': weights,
+        },
+        jnp.asarray(True),
+    )
+
+    chex.assert_trees_all_close(
+        nnx.state(accumulated_model, nnx.Param),
+        nnx.state(reference_model, nnx.Param),
+        rtol=1e-6,
+        atol=1e-6,
+    )
+
+  def test_loss_output_metrics_use_global_denominator(self):
+    inputs = jnp.array(
+        [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [2.0, -1.0]],
+        dtype=jnp.float32,
+    )
+    targets = jnp.array([[0.5], [-0.5], [1.0], [-1.0]], dtype=jnp.float32)
+    weights = jnp.array([0.1, 0.0, 0.2, 0.3], dtype=jnp.float32)
+
+    def weighted_loss_sum(model, x, y, example_weights):
+      per_example = jnp.sum(jnp.square(model(x) - y), axis=-1)
+      return jnp.sum(per_example * example_weights)
+
+    def loss_fn(model, x, y, example_weights):
+      metric = utils.WeightedMetric(
+          weighted_loss_sum(model, x, y, example_weights),
+          jnp.sum(example_weights),
+      )
+      return utils.LossOutput(
+          primary_loss=metric,
+          aux_metrics={'weighted_mse': metric},
+      )
+
+    reference_model = nnx.Linear(2, 1, rngs=nnx.Rngs(0))
+    expected = weighted_loss_sum(
+        reference_model, inputs, targets, weights
+    ) / jnp.sum(weights)
+    model = nnx.Linear(2, 1, rngs=nnx.Rngs(0))
+    trainer = peft_trainer.PeftTrainer(
+        model,
+        optax.sgd(0.1),
+        peft_trainer.TrainingConfig(
+            eval_every_n_steps=100,
+            max_steps=1,
+            gradient_accumulation_steps=2,
+        ),
+    ).with_loss_fn(loss_fn)
+    train_batches = [
+        {
+            'x': inputs[start : start + 2],
+            'y': targets[start : start + 2],
+            'example_weights': weights[start : start + 2],
+        }
+        for start in (0, 2)
+    ]
+
+    trainer.train(train_batches)
+
+    np.testing.assert_allclose(
+        trainer.metrics_logger.get_metric('', 'loss', 'train'),
+        expected,
+        rtol=1e-6,
+    )
+    np.testing.assert_allclose(
+        trainer.metrics_logger.get_metric('', 'weighted_mse', 'train'),
+        expected,
+        rtol=1e-6,
+    )
+
   def test_stream_train_step_returns_raw_weighted_metric_aux(self):
     # Since _compute_legacy_aux was dropped, the (stream) _train_step returns
     # the raw aux with WeightedMetric preserved, so the metric ops can reduce
@@ -801,15 +1048,24 @@ def loss_fn(
         trainer.model, trainer.optimizer, acc, self.train_ds[0], jnp.array(True)
     )
 
+    self.assertIsInstance(aux, utils.LossOutput)
+    aux_metrics = aux.aux_metrics
     # aux values stay raw WeightedMetric (not pre-divided to scalars).
-    self.assertIsInstance(aux['foo'], utils.WeightedMetric)
-    self.assertIsInstance(aux['bar'], utils.WeightedMetric)
+    self.assertIsInstance(aux_metrics['foo'], utils.WeightedMetric)
+    self.assertIsInstance(aux_metrics['bar'], utils.WeightedMetric)
     # Metric ops reduce the raw WeightedMetric correctly across micro-batches.
     self.assertAlmostEqual(
-        float(common.mean_of_means([aux['foo'], aux['foo']])), 2.0, places=5
+        float(common.mean_of_means([aux_metrics['foo'], aux_metrics['foo']])),
+        2.0,
+        places=5,
     )  # mean(10/5, 10/5)
     self.assertAlmostEqual(
-        float(common.global_weighted_mean([aux['bar'], aux['bar']])), 3.0,
+        float(
+            common.global_weighted_mean(
+                [aux_metrics['bar'], aux_metrics['bar']]
+            )
+        ),
+        3.0,
         places=5,
     )  # (6+6)/(2+2)
 
@@ -891,6 +1147,7 @@ def test_default_mode_averages_grads(self):
       dict(testcase_name='equal_denoms', denoms=(4.0, 4.0, 4.0, 4.0)),
       dict(testcase_name='varying_denoms', denoms=(1.0, 7.0, 3.0, 5.0)),
       dict(testcase_name='extreme_variance', denoms=(1.0, 1.0, 100.0, 1.0)),
+      dict(testcase_name='fractional_denoms', denoms=(0.05, 0.1, 0.2, 0.25)),
   )
   def test_explicit_denom_matches_single_step_baseline(self, denoms):
     """Passing explicit denom matches the equivalent single-step batch.
@@ -1308,43 +1565,85 @@ def step(model, optimizer, acc, is_update_step):
     for d in float_dtypes:
       self.assertEqual(d, jnp.bfloat16)
 
-  def test_peft_trainer_promotes_bf16_opt_state_floats_to_float32(self):
-    """`PeftTrainer.__init__` casts float opt_state leaves to float32.
-
-    `optax.adam` / `optax.adamw` promote their floating-point moments
-    (`mu`, `nu`) to float32 inside `update` whenever the learning rate is
-    a float32 tracer (as produced by `optax.inject_hyperparams`). This test
-    verifies that the trainer casts these to float32 in-place during init.
-    """
+  @parameterized.product(
+      optimizer_kind=('direct', 'injected'),
+      gradient_accumulation_steps=(1, 2),
+  )
+  def test_peft_trainer_preserves_bf16_optimizer_state_dtypes_under_jit(
+      self, optimizer_kind, gradient_accumulation_steps
+  ):
+    """Jitted update/skip branches preserve each optimizer's state contract."""
     rngs = nnx.Rngs(0)
-    model = tc.ToyTransformer(config=tc.ModelConfig(), rngs=rngs)
-    bf16_state = jax.tree.map(
-        lambda x: x.astype(jnp.bfloat16)
-        if jnp.issubdtype(x.dtype, jnp.floating)
-        else x,
-        nnx.state(model, nnx.Param),
+    model = nnx.Linear(
+        in_features=4, out_features=2, rngs=rngs, param_dtype=jnp.bfloat16
     )
-    nnx.update(model, bf16_state)
-
-    tx = optax.inject_hyperparams(optax.adamw, hyperparam_dtype=jnp.float32)(
-        learning_rate=1e-3
+    if optimizer_kind == 'direct':
+      tx = optax.adamw(
+          lambda _: jnp.asarray(0.1, dtype=jnp.float32),
+          mu_dtype=jnp.bfloat16,
+      )
+    else:
+      tx = optax.inject_hyperparams(optax.adamw, hyperparam_dtype=jnp.float32)(
+          learning_rate=0.1
+      )
+    config = peft_trainer.TrainingConfig(
+        eval_every_n_steps=100,
+        max_steps=1,
+        gradient_accumulation_steps=gradient_accumulation_steps,
     )
-    config = peft_trainer.TrainingConfig(eval_every_n_steps=100, max_steps=1)
     trainer = peft_trainer.PeftTrainer(model, tx, config)
+    trainer.with_loss_fn(lambda m, x, y: jnp.sum((m(x) - y) ** 2))
 
-    opt_state_dtypes = jax.tree_util.tree_leaves(
-        jax.tree_util.tree_map(
-            lambda v: v[...].dtype,
-            nnx.state(trainer.optimizer, nnx.optimizer.OptState),
-            is_leaf=lambda x: isinstance(x, nnx.Variable),
+    def opt_state_dtypes():
+      return jax.tree_util.tree_map(
+          lambda v: v[...].dtype,
+          nnx.state(trainer.optimizer, nnx.optimizer.OptState),
+          is_leaf=lambda x: isinstance(x, nnx.Variable),
+      )
+
+    initial_opt_state_dtypes = opt_state_dtypes()
+    initial_params = jax.tree_util.tree_map(
+        lambda v: jnp.copy(v[...]),
+        nnx.state(model, nnx.Param),
+        is_leaf=lambda x: isinstance(x, nnx.Variable),
+    )
+    train_step, _ = trainer.jit_train_and_eval_step()
+    inputs = {
+        'x': jnp.ones((2, 4), dtype=jnp.bfloat16),
+        'y': jnp.ones((2, 2), dtype=jnp.bfloat16),
+    }
+
+    for micro_step in range(gradient_accumulation_steps):
+      is_update_step = micro_step == gradient_accumulation_steps - 1
+      _, _, grad_norm = train_step(
+          inputs, is_update_step=jnp.asarray(is_update_step)
+      )
+      self.assertEqual(grad_norm.dtype, jnp.float32)
+      if not is_update_step:
+        for before, after in zip(
+            jax.tree_util.tree_leaves(initial_params),
+            jax.tree_util.tree_leaves(
+                nnx.state(model, nnx.Param),
+                is_leaf=lambda x: isinstance(x, nnx.Variable),
+            ),
+        ):
+          np.testing.assert_array_equal(before, after[...])
+
+    self.assertEqual(opt_state_dtypes(), initial_opt_state_dtypes)
+    updated_params = nnx.state(model, nnx.Param)
+    self.assertTrue(
+        any(
+            not np.array_equal(before, after[...])
+            for before, after in zip(
+                jax.tree_util.tree_leaves(initial_params),
+                jax.tree_util.tree_leaves(
+                    updated_params,
+                    is_leaf=lambda x: isinstance(x, nnx.Variable),
+                ),
+            )
         )
     )
-    float_dtypes = [
-        d for d in opt_state_dtypes if jnp.issubdtype(d, jnp.floating)
-    ]
-    self.assertNotEmpty(float_dtypes)
-    for d in float_dtypes:
-      self.assertEqual(d, jnp.float32)
+    self.assertEqual(float(trainer.grad_accumulator.denom[...]), 0.0)
 
 
 if __name__ == '__main__':
diff --git a/tunix/diffusion/__init__.py b/tunix/diffusion/__init__.py
new file mode 100644
index 000000000..d6d7357c9
--- /dev/null
+++ b/tunix/diffusion/__init__.py
@@ -0,0 +1,29 @@
+# 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.
+
+"""Framework-neutral contracts for diffusion training."""
+
+# pylint: disable=g-importing-member
+
+from tunix.diffusion.interfaces import compute_diffusion_logits
+from tunix.diffusion.interfaces import DiffusionBatchAdapter
+from tunix.diffusion.interfaces import DiffusionLogitsFn
+from tunix.diffusion.types import DiffusionTokenBatch
+
+__all__ = [
+    "DiffusionBatchAdapter",
+    "DiffusionLogitsFn",
+    "DiffusionTokenBatch",
+    "compute_diffusion_logits",
+]
diff --git a/tunix/diffusion/interfaces.py b/tunix/diffusion/interfaces.py
new file mode 100644
index 000000000..7e1d3c3a9
--- /dev/null
+++ b/tunix/diffusion/interfaces.py
@@ -0,0 +1,53 @@
+# 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.
+
+"""Callable interfaces for framework-neutral diffusion integrations."""
+
+from typing import Protocol, TypeVar
+
+from flax import nnx
+from tunix.diffusion import types
+
+RawBatchT_contra = TypeVar("RawBatchT_contra", contravariant=True)
+
+
+class DiffusionBatchAdapter(Protocol[RawBatchT_contra]):
+  """Converts an external batch to Tunix's canonical diffusion batch."""
+
+  def __call__(self, batch: RawBatchT_contra, /) -> types.DiffusionTokenBatch:
+    ...
+
+
+class DiffusionLogitsFn(Protocol):
+  """Computes target-aligned token logits for a diffusion model."""
+
+  def __call__(
+      self,
+      model: nnx.Module,
+      model_inputs: types.ModelInputs,
+      /,
+  ) -> types.Array:
+    ...
+
+
+def compute_diffusion_logits(
+    model: nnx.Module,
+    batch: types.DiffusionTokenBatch,
+    logits_fn: DiffusionLogitsFn,
+) -> types.Array:
+  """Computes diffusion logits and validates target alignment."""
+
+  batch.validate()
+  logits = logits_fn(model, batch.model_inputs)
+  return types.validate_diffusion_logits(batch, logits)
diff --git a/tunix/diffusion/types.py b/tunix/diffusion/types.py
new file mode 100644
index 000000000..19e563f54
--- /dev/null
+++ b/tunix/diffusion/types.py
@@ -0,0 +1,200 @@
+# 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.
+
+"""Data types shared by diffusion training objectives."""
+
+from collections.abc import Mapping
+from typing import Any, TypeAlias
+
+import flax
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+Array: TypeAlias = jax.Array | np.ndarray
+ModelInputs: TypeAlias = Mapping[str, Any]
+
+
+def _shape_and_dtype(name: str, value: Any) -> tuple[tuple[int, ...], Any]:
+  if not isinstance(value, (jax.Array, jax.core.Tracer, np.ndarray)):
+    raise TypeError(f"{name} must be a JAX or NumPy array")
+  return tuple(value.shape), value.dtype
+
+
+def _is_loss_weight_dtype(dtype: Any) -> bool:
+  return (
+      jnp.issubdtype(dtype, jnp.bool_)
+      or jnp.issubdtype(dtype, jnp.integer)
+      or jnp.issubdtype(dtype, jnp.floating)
+  )
+
+
+def _concrete_numpy(value: Any) -> np.ndarray | None:
+  """Returns eager values without performing host conversion while tracing."""
+
+  if isinstance(value, jax.core.Tracer):
+    return None
+  if isinstance(value, jax.Array) and not value.is_fully_addressable:
+    return None
+  return np.asarray(value)
+
+
+@flax.struct.dataclass(frozen=True)
+class DiffusionTokenBatch:
+  """Canonical target-aligned batch for diffusion training.
+
+  Every array in ``model_inputs`` must be batch-major. Model-static values and
+  non-batch tensors belong in the scoring callable so that Tunix can shard the
+  complete batch pytree consistently along its leading dimension.
+
+  Attributes:
+    model_inputs: Model-specific, batch-major array pytree consumed by a
+      ``DiffusionLogitsFn``.
+    target_ids: Target token IDs with shape ``[batch, length]``.
+    loss_weights: Per-target weights with the same shape as ``target_ids``.
+  """
+
+  model_inputs: ModelInputs
+  target_ids: Array
+  loss_weights: Array
+
+  @classmethod
+  def create(
+      cls,
+      *,
+      model_inputs: ModelInputs,
+      target_ids: Array,
+      loss_weights: Array,
+  ) -> "DiffusionTokenBatch":
+    """Constructs and validates a diffusion batch."""
+
+    return cls(
+        model_inputs=model_inputs,
+        target_ids=target_ids,
+        loss_weights=loss_weights,
+    ).validate()
+
+  def validate(self) -> "DiffusionTokenBatch":
+    """Validates the batch's static shape and dtype contract."""
+
+    if not isinstance(self.model_inputs, Mapping):
+      raise TypeError("model_inputs must be a mapping")
+
+    target_shape, target_dtype = _shape_and_dtype("target_ids", self.target_ids)
+    if len(target_shape) != 2:
+      raise ValueError(
+          f"target_ids must have shape [batch, length]; received {target_shape}"
+      )
+    if not jnp.issubdtype(target_dtype, jnp.integer):
+      raise TypeError(
+          f"target_ids must have an integer dtype; received {target_dtype}"
+      )
+
+    weight_shape, weight_dtype = _shape_and_dtype(
+        "loss_weights", self.loss_weights
+    )
+    if weight_shape != target_shape:
+      raise ValueError(
+          "loss_weights must match target_ids shape; "
+          f"received {weight_shape} and {target_shape}"
+      )
+    if not _is_loss_weight_dtype(weight_dtype):
+      raise TypeError(
+          "loss_weights must have a real numeric or boolean dtype; "
+          f"received {weight_dtype}"
+      )
+    concrete_weights = _concrete_numpy(self.loss_weights)
+    if concrete_weights is not None:
+      if not np.all(np.isfinite(concrete_weights)):
+        raise ValueError("loss_weights must contain only finite values")
+      if np.any(concrete_weights < 0):
+        raise ValueError("loss_weights must be nonnegative")
+      concrete_targets = _concrete_numpy(self.target_ids)
+      if concrete_targets is not None and np.any(
+          concrete_targets[concrete_weights > 0] < 0
+      ):
+        raise ValueError("active target_ids must be nonnegative")
+
+    model_input_leaves = jax.tree.leaves(self.model_inputs)
+    if not model_input_leaves:
+      raise ValueError("model_inputs must contain at least one array")
+    batch_size = target_shape[0]
+    for index, leaf in enumerate(model_input_leaves):
+      leaf_shape, leaf_dtype = _shape_and_dtype(
+          f"model_inputs leaf {index}", leaf
+      )
+      if not leaf_shape:
+        raise ValueError(
+            f"model_inputs leaf {index} must be batch-major, not scalar"
+        )
+      if leaf_shape[0] != batch_size:
+        raise ValueError(
+            f"model_inputs leaf {index} has batch size {leaf_shape[0]}; "
+            f"expected {batch_size}"
+        )
+      if not _is_loss_weight_dtype(leaf_dtype):
+        raise TypeError(
+            f"model_inputs leaf {index} must have a real numeric or boolean "
+            f"dtype; received {leaf_dtype}"
+        )
+    return self
+
+
+def validate_diffusion_logits(
+    batch: DiffusionTokenBatch, logits: Array
+) -> Array:
+  """Validates and returns target-aligned token logits.
+
+  The checks use static shape and dtype metadata, so this function is safe to
+  invoke while tracing a JAX computation.
+
+  Args:
+    batch: The batch whose targets the scores must align with.
+    logits: Floating-point token logits with shape ``[batch, length, vocab]``.
+
+  Returns:
+    ``logits`` unchanged.
+
+  Raises:
+    TypeError: If ``logits`` is not an array or has a non-floating dtype.
+    ValueError: If ``logits`` is not target-aligned, has an empty vocabulary,
+      or an active target ID falls outside that vocabulary.
+  """
+
+  logit_shape, logit_dtype = _shape_and_dtype("logits", logits)
+  if len(logit_shape) != 3:
+    raise ValueError(
+        f"logits must have shape [batch, length, vocab]; received {logit_shape}"
+    )
+  if logit_shape[:2] != tuple(batch.target_ids.shape):
+    raise ValueError(
+        "logits must align with target_ids on [batch, length]; "
+        f"received {logit_shape[:2]} and {tuple(batch.target_ids.shape)}"
+    )
+  if logit_shape[2] < 1:
+    raise ValueError("logits must have a non-empty vocabulary dimension")
+  if not jnp.issubdtype(logit_dtype, jnp.floating):
+    raise TypeError(
+        f"logits must have a floating-point dtype; received {logit_dtype}"
+    )
+  concrete_targets = _concrete_numpy(batch.target_ids)
+  concrete_weights = _concrete_numpy(batch.loss_weights)
+  if concrete_targets is not None and concrete_weights is not None:
+    active_targets = concrete_targets[concrete_weights > 0]
+    if np.any(active_targets >= logit_shape[2]):
+      raise ValueError(
+          "active target_ids must be smaller than vocabulary size"
+          f" {logit_shape[2]}"
+      )
+  return logits
diff --git a/tunix/distillation/diffusion.py b/tunix/distillation/diffusion.py
new file mode 100644
index 000000000..420785afc
--- /dev/null
+++ b/tunix/distillation/diffusion.py
@@ -0,0 +1,74 @@
+# 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.
+
+"""External-teacher batch contracts for diffusion distillation."""
+
+from typing import Protocol, TypeVar
+
+import flax
+from tunix.diffusion import types as diffusion_types
+
+RawBatchT_contra = TypeVar("RawBatchT_contra", contravariant=True)
+
+
+@flax.struct.dataclass(frozen=True)
+class DiffusionDistillationBatch:
+  """Student inputs and external teacher logits for diffusion distillation.
+
+  An upstream model-aware integration is responsible for generating an
+  on-policy student rollout and scoring that same rollout with the teacher.
+  Tunix only consumes the resulting target-aligned tensors, so it does not own
+  model-specific rollout, corruption, or checkpoint behavior.
+
+  Attributes:
+    student_batch: Inputs, hard targets, and per-target loss weights used to
+      score the student.
+    teacher_logits: Immutable teacher logits with shape ``[batch, length,
+      vocab]`` aligned with ``student_batch.target_ids``.
+  """
+
+  student_batch: diffusion_types.DiffusionTokenBatch
+  teacher_logits: diffusion_types.Array
+
+  @classmethod
+  def create(
+      cls,
+      *,
+      student_batch: diffusion_types.DiffusionTokenBatch,
+      teacher_logits: diffusion_types.Array,
+  ) -> "DiffusionDistillationBatch":
+    """Constructs and validates a diffusion distillation batch."""
+
+    return cls(
+        student_batch=student_batch,
+        teacher_logits=teacher_logits,
+    ).validate()
+
+  def validate(self) -> "DiffusionDistillationBatch":
+    """Validates the student batch and target alignment of teacher logits."""
+
+    if not isinstance(self.student_batch, diffusion_types.DiffusionTokenBatch):
+      raise TypeError("student_batch must be a DiffusionTokenBatch")
+    self.student_batch.validate()
+    diffusion_types.validate_diffusion_logits(
+        self.student_batch, self.teacher_logits
+    )
+    return self
+
+
+class PreparedDiffusionDistillationBatchAdapter(Protocol[RawBatchT_contra]):
+  """Converts a freshly prepared external rollout to the OPD contract."""
+
+  def __call__(self, batch: RawBatchT_contra, /) -> DiffusionDistillationBatch:
+    ...
diff --git a/tunix/distillation/diffusion_opd.py b/tunix/distillation/diffusion_opd.py
new file mode 100644
index 000000000..ce5fc8ba2
--- /dev/null
+++ b/tunix/distillation/diffusion_opd.py
@@ -0,0 +1,209 @@
+# 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.
+
+"""On-policy distillation over target-aligned diffusion scores."""
+
+import math
+import numbers
+from typing import Any, TypeVar
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+from tunix.diffusion import interfaces as diffusion_interfaces
+from tunix.distillation import diffusion
+from tunix.sft import peft_trainer
+from tunix.sft import utils as sft_utils
+
+RawBatchT = TypeVar("RawBatchT")
+TrainerT = TypeVar("TrainerT", bound=peft_trainer.PeftTrainer)
+
+
+def _validate_loss_configuration(
+    temperature: float,
+    soft_loss_weight: float,
+    hard_loss_weight: float,
+) -> None:
+  """Validates static diffusion distillation loss settings."""
+
+  values = {
+      "temperature": temperature,
+      "soft_loss_weight": soft_loss_weight,
+      "hard_loss_weight": hard_loss_weight,
+  }
+  for name, value in values.items():
+    if (
+        isinstance(value, bool)
+        or not isinstance(value, numbers.Real)
+        or not math.isfinite(float(value))
+    ):
+      raise ValueError(f"{name} must be a finite real scalar; received {value}")
+  if temperature <= 0:
+    raise ValueError(f"temperature must be positive; received {temperature}")
+  if soft_loss_weight < 0 or hard_loss_weight < 0:
+    raise ValueError(
+        "soft_loss_weight and hard_loss_weight must be nonnegative"
+    )
+  if soft_loss_weight == 0 and hard_loss_weight == 0:
+    raise ValueError("at least one loss weight must be positive")
+
+
+def _validate_active_teacher_logits(
+    batch: diffusion.DiffusionDistillationBatch,
+) -> None:
+  """Rejects invalid eager teacher logits without converting JAX tracers."""
+
+  if isinstance(batch.teacher_logits, jax.core.Tracer) or isinstance(
+      batch.student_batch.loss_weights, jax.core.Tracer
+  ):
+    return
+  teacher_logits = np.asarray(batch.teacher_logits)
+  active_targets = np.asarray(batch.student_batch.loss_weights) > 0
+  if not np.all(np.isfinite(teacher_logits[active_targets])):
+    raise ValueError("active teacher_logits must contain only finite values")
+
+
+def diffusion_opd_loss_fn(
+    student_model: nnx.Module,
+    batch: diffusion.DiffusionDistillationBatch,
+    student_logits_fn: diffusion_interfaces.DiffusionLogitsFn,
+    *,
+    temperature: float = 1.0,
+    soft_loss_weight: float = 1.0,
+    hard_loss_weight: float = 0.0,
+) -> sft_utils.LossOutput:
+  """Computes weighted on-policy distillation loss for diffusion scores.
+
+  Soft loss is the forward KL ``KL(teacher || student)`` over the same
+  target-aligned student rollout, with standard temperature-squared scaling.
+  Hard loss is optional target-aligned cross-entropy against ``target_ids``.
+  Teacher logits are always stop-gradient.
+
+  Args:
+    student_model: The trainable student model.
+    batch: Target-aligned student inputs and external teacher logits.
+    student_logits_fn: Callable that returns student logits with shape
+      ``[batch, length, vocab]``.
+    temperature: Positive softmax temperature for the KL term.
+    soft_loss_weight: Nonnegative multiplier for forward KL.
+    hard_loss_weight: Nonnegative multiplier for hard-target cross-entropy.
+
+  Returns:
+    An unreduced weighted loss and token-weighted component metrics.
+  """
+
+  _validate_loss_configuration(temperature, soft_loss_weight, hard_loss_weight)
+  batch.validate()
+  if soft_loss_weight:
+    _validate_active_teacher_logits(batch)
+
+  student_logits = diffusion_interfaces.compute_diffusion_logits(
+      student_model, batch.student_batch, student_logits_fn
+  )
+  if tuple(student_logits.shape) != tuple(batch.teacher_logits.shape):
+    raise ValueError(
+        "student and teacher logits must have identical shapes; received "
+        f"{tuple(student_logits.shape)} and {tuple(batch.teacher_logits.shape)}"
+    )
+
+  weights = jnp.asarray(batch.student_batch.loss_weights, dtype=jnp.float32)
+  active_targets = weights != 0
+  student_logits = jnp.asarray(student_logits, dtype=jnp.float32)
+  teacher_logits = jax.lax.stop_gradient(
+      jnp.asarray(batch.teacher_logits, dtype=jnp.float32)
+  )
+  student_logits = jnp.where(active_targets[..., None], student_logits, 0.0)
+  teacher_logits = jnp.where(active_targets[..., None], teacher_logits, 0.0)
+
+  if soft_loss_weight:
+    student_log_probs = jax.nn.log_softmax(
+        student_logits / temperature, axis=-1
+    )
+    teacher_probs = jax.nn.softmax(teacher_logits / temperature, axis=-1)
+    soft_token_losses = optax.kl_divergence(
+        student_log_probs, teacher_probs
+    ) * (temperature**2)
+  else:
+    soft_token_losses = jnp.zeros_like(weights)
+
+  if hard_loss_weight:
+    targets = jnp.asarray(batch.student_batch.target_ids)
+    targets = jnp.where(active_targets, targets, 0)
+    hard_token_losses = optax.softmax_cross_entropy_with_integer_labels(
+        logits=student_logits,
+        labels=targets,
+    )
+  else:
+    hard_token_losses = jnp.zeros_like(weights)
+
+  soft_sum = jnp.sum(soft_token_losses * weights, dtype=jnp.float32)
+  hard_sum = jnp.sum(hard_token_losses * weights, dtype=jnp.float32)
+  total_sum = (
+      jnp.asarray(soft_loss_weight, dtype=jnp.float32) * soft_sum
+      + jnp.asarray(hard_loss_weight, dtype=jnp.float32) * hard_sum
+  )
+  weight_sum = jnp.sum(weights, dtype=jnp.float32)
+
+  return sft_utils.LossOutput(
+      primary_loss=sft_utils.WeightedMetric(total_sum, weight_sum),
+      aux_metrics={
+          "distill/soft_loss": sft_utils.WeightedMetric(soft_sum, weight_sum),
+          "distill/hard_loss": sft_utils.WeightedMetric(hard_sum, weight_sum),
+      },
+  )
+
+
+def configure_prepared_diffusion_opd(
+    trainer: TrainerT,
+    batch_adapter: diffusion.PreparedDiffusionDistillationBatchAdapter[
+        RawBatchT
+    ],
+    student_logits_fn: diffusion_interfaces.DiffusionLogitsFn,
+    *,
+    temperature: float = 1.0,
+    soft_loss_weight: float = 1.0,
+    hard_loss_weight: float = 0.0,
+) -> TrainerT:
+  """Configures a ``PeftTrainer`` for prepared external-teacher diffusion OPD.
+
+  The raw batch must contain a fresh rollout from the current student and
+  teacher logits for that exact rollout. This function does not perform or
+  verify on-policy generation; a model-aware integration must prepare the
+  batch outside the jitted loss step. It only binds the prepared adapter and
+  objective to Tunix's stable trainer extension points.
+  """
+
+  _validate_loss_configuration(temperature, soft_loss_weight, hard_loss_weight)
+
+  def gen_model_input_fn(raw_batch: RawBatchT) -> dict[str, Any]:
+    return {"batch": batch_adapter(raw_batch)}
+
+  def loss_fn(
+      student_model: nnx.Module,
+      batch: diffusion.DiffusionDistillationBatch,
+  ) -> sft_utils.LossOutput:
+    return diffusion_opd_loss_fn(
+        student_model,
+        batch,
+        student_logits_fn,
+        temperature=temperature,
+        soft_loss_weight=soft_loss_weight,
+        hard_loss_weight=hard_loss_weight,
+    )
+
+  trainer.with_gen_model_input_fn(gen_model_input_fn)
+  trainer.with_loss_fn(loss_fn)
+  return trainer
diff --git a/tunix/experimental/train/peft_trainer_v2.py b/tunix/experimental/train/peft_trainer_v2.py
index a61ca934f..2379ce792 100644
--- a/tunix/experimental/train/peft_trainer_v2.py
+++ b/tunix/experimental/train/peft_trainer_v2.py
@@ -33,6 +33,7 @@
 import optax
 import orbax.checkpoint as ocp
 from tunix.experimental.common import datatypes
+from tunix.experimental.metrics import metrics as exp_metrics
 from tunix.experimental.train import abstract_trainer
 from tunix.perf import metrics as perf_metrics
 from tunix.perf import trace as perf_trace
@@ -45,7 +46,6 @@
 from tunix.sft import profiler
 from tunix.sft import progress_bar
 from tunix.sft import sharding_utils
-from tunix.experimental.metrics import metrics as exp_metrics
 from tunix.sft import utils
 from typing_extensions import override
 
@@ -433,7 +433,9 @@ def with_gen_model_input_fn(
       PeftTrainer.
     """
     self.clear_jit_cache()
-    self.gen_model_input_fn = gen_model_input_fn  # pyrefly: ignore[bad-assignment]
+    self.gen_model_input_fn = (
+        gen_model_input_fn  # pyrefly: ignore[bad-assignment]
+    )
     return self
 
   def _fwd_bwd_step(
@@ -683,7 +685,9 @@ def _log_metrics(
   ):
     """Logs the metrics to the metrics logger and console."""
     perplexity = np.exp(jax.device_get(loss))
-    self.metrics_logger.log(self.metrics_prefix, "loss", loss, self._mode, step)  # pyrefly: ignore[missing-attribute]
+    self.metrics_logger.log(
+        self.metrics_prefix, "loss", loss, self._mode, step
+    )  # pyrefly: ignore[missing-attribute]
     self.metrics_logger.log(  # pyrefly: ignore[missing-attribute]
         self.metrics_prefix, "perplexity", perplexity, self._mode, step
     )
@@ -705,7 +709,9 @@ def _log_metrics(
           perplexity,
       )
     for k, v in (additional_metrics or {}).items():
-      self.metrics_logger.log(self.metrics_prefix, k, v, self._mode, step)  # pyrefly: ignore[missing-attribute]
+      self.metrics_logger.log(
+          self.metrics_prefix, k, v, self._mode, step
+      )  # pyrefly: ignore[missing-attribute]
 
   def _buffer_metrics(
       self,
@@ -773,9 +779,7 @@ def _to_np_array(v):
     weighted_metrics = {}
     scalar_metrics = {"loss": loss}
     for k, val in additional_metrics.items():
-      if isinstance(
-          val, (utils.WeightedMetric, exp_metrics.WeightedMetric)
-      ):
+      if isinstance(val, (utils.WeightedMetric, exp_metrics.WeightedMetric)):
         weighted_metrics[k] = val
       else:
         scalar_metrics[k] = val
@@ -950,8 +954,8 @@ def train(
         f"Training with mesh: {pxla.thread_resources.env.physical_mesh}",
         1,
     )
-    fwd_bwd_step, _, _ = (
-        self.jit_fwd_bwd_update_and_eval_step(skip_jit, cache_nnx_graph)
+    fwd_bwd_step, _, _ = self.jit_fwd_bwd_update_and_eval_step(
+        skip_jit, cache_nnx_graph
     )
     if not skip_jit:
       cache_size = fwd_bwd_step.func.jitted_fn._cache_size()  # pytype: disable=attribute-error
@@ -1109,6 +1113,7 @@ def _save_last_checkpoint(self):
           self.model,
           self.optimizer,
           save_only_lora_params=self._lora_enabled,
+          custom_metadata=self.custom_checkpoint_metadata(),
           force=True,
       )
 
diff --git a/tunix/rl/diffusion.py b/tunix/rl/diffusion.py
new file mode 100644
index 000000000..33384d89c
--- /dev/null
+++ b/tunix/rl/diffusion.py
@@ -0,0 +1,103 @@
+# 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.
+
+"""Policy scoring for prepared, target-aligned diffusion rollouts."""
+
+import math
+import numbers
+
+from flax import nnx
+import jax
+import jax.numpy as jnp
+from tunix.diffusion import interfaces as diffusion_interfaces
+from tunix.diffusion import types as diffusion_types
+
+
+def _validate_temperature(temperature: float) -> None:
+  if (
+      isinstance(temperature, bool)
+      or not isinstance(temperature, numbers.Real)
+      or not math.isfinite(float(temperature))
+      or temperature <= 0
+  ):
+    raise ValueError(
+        "temperature must be a finite positive real scalar; received"
+        f" {temperature}"
+    )
+
+
+def compute_diffusion_per_token_logps(
+    model: nnx.Module,
+    batch: diffusion_types.DiffusionTokenBatch,
+    logits_fn: diffusion_interfaces.DiffusionLogitsFn,
+    *,
+    temperature: float = 1.0,
+    stop_gradient: bool = True,
+    return_entropy: bool = False,
+) -> jax.Array | tuple[jax.Array, jax.Array]:
+  """Scores a prepared diffusion rollout without an autoregressive shift.
+
+  The caller owns model-specific rollout and batch construction. ``batch``
+  must describe the exact target-aligned canvas used to score the completion.
+  Positions with zero ``loss_weights`` are inactive: their logits and target
+  IDs are sanitized before probability operations, then their returned scores
+  are zeroed. This permits padding to carry sentinel IDs or non-finite logits
+  without contaminating active scores or gradients.
+
+  Args:
+    model: Policy model to score.
+    batch: Prepared target-aligned model inputs, token IDs, and active weights.
+    logits_fn: Model-specific callable returning ``[batch, length, vocab]``
+      target-aligned logits.
+    temperature: Positive sampling temperature used by the rollout policy.
+    stop_gradient: Whether returned scores should be detached from the policy.
+    return_entropy: Whether to also return per-token categorical entropy.
+
+  Returns:
+    Per-token log probabilities with the same shape as ``batch.target_ids``.
+    When ``return_entropy`` is true, also returns entropy with the same shape.
+    Inactive positions are zero in every returned array.
+  """
+
+  _validate_temperature(temperature)
+  logits = diffusion_interfaces.compute_diffusion_logits(
+      model, batch, logits_fn
+  )
+
+  active = jnp.asarray(batch.loss_weights) != 0
+  logits = jnp.asarray(logits, dtype=jnp.float32)
+  logits = jnp.where(active[..., None], logits, 0.0)
+  targets = jnp.asarray(batch.target_ids)
+  targets = jnp.where(active, targets, 0)
+
+  log_probs = jax.nn.log_softmax(logits / temperature, axis=-1)
+  per_token_logps = jnp.take_along_axis(log_probs, targets[..., None], axis=-1)[
+      ..., 0
+  ]
+  per_token_logps = jnp.where(active, per_token_logps, 0.0)
+
+  if return_entropy:
+    probs = jnp.exp(log_probs)
+    entropy_terms = jnp.where(probs > 0, probs * log_probs, 0.0)
+    entropy = -jnp.sum(entropy_terms, axis=-1)
+    entropy = jnp.where(active, entropy, 0.0)
+
+  if stop_gradient:
+    per_token_logps = jax.lax.stop_gradient(per_token_logps)
+    if return_entropy:
+      entropy = jax.lax.stop_gradient(entropy)
+
+  if return_entropy:
+    return per_token_logps, entropy
+  return per_token_logps
diff --git a/tunix/rl/rl_cluster.py b/tunix/rl/rl_cluster.py
index 7a5118c1c..561a57e2e 100644
--- a/tunix/rl/rl_cluster.py
+++ b/tunix/rl/rl_cluster.py
@@ -38,6 +38,7 @@
 import jaxtyping
 import numpy as np
 import optax
+from tunix.diffusion import types as diffusion_types
 from tunix.generate import tokenizer_adapter
 # Internal placeholder for sglang_jax rollout worker stub, don't change this line.
 from tunix.perf import metrics as perf_metrics
@@ -61,6 +62,60 @@
 MetricsBuffer = perf_metrics.MetricsBuffer
 
 
+def _merge_diffusion_batches(
+    outputs: list[base_rollout.RolloutOutput],
+) -> diffusion_types.DiffusionTokenBatch | None:
+  """Merges optional prepared diffusion metadata across rollout microbatches."""
+
+  batches = [output.diffusion_batch for output in outputs]
+  present = [batch is not None for batch in batches]
+  if any(present) and not all(present):
+    raise ValueError(
+        "rollout microbatches must either all provide diffusion_batch or all "
+        "omit it"
+    )
+  if not any(present):
+    return None
+
+  diffusion_batches = []
+  for index, (output, batch) in enumerate(zip(outputs, batches)):
+    if batch is None:
+      continue
+    batch = batch.validate()
+    if batch.target_ids.shape[0] != len(output.text):
+      raise ValueError(
+          "diffusion_batch batch size must match its rollout output; "
+          f"microbatch {index} has {batch.target_ids.shape[0]} prepared "
+          f"examples and {len(output.text)} generated examples"
+      )
+    diffusion_batches.append(batch)
+  model_input_structure = jax.tree.structure(diffusion_batches[0].model_inputs)
+  if any(
+      jax.tree.structure(batch.model_inputs) != model_input_structure
+      for batch in diffusion_batches[1:]
+  ):
+    raise ValueError(
+        "diffusion_batch model_inputs must have the same pytree structure "
+        "across rollout microbatches"
+    )
+
+  def concatenate(*arrays):
+    return jnp.concatenate([jnp.asarray(array) for array in arrays], axis=0)
+
+  return diffusion_types.DiffusionTokenBatch.create(
+      model_inputs=jax.tree.map(
+          concatenate,
+          *(batch.model_inputs for batch in diffusion_batches),
+      ),
+      target_ids=concatenate(
+          *(batch.target_ids for batch in diffusion_batches)
+      ),
+      loss_weights=concatenate(
+          *(batch.loss_weights for batch in diffusion_batches)
+      ),
+  )
+
+
 class Mode(enum.Enum):
   """Mode of RolloutConfig."""
 
@@ -980,6 +1035,7 @@ def generate(
             [out.left_padded_prompt_tokens for out in outputs], axis=0
         ),
         logprobs=logprobs,
+        diffusion_batch=_merge_diffusion_batches(outputs),
     )
 
   def get_ref_per_token_logps(
diff --git a/tunix/rl/rollout/base_rollout.py b/tunix/rl/rollout/base_rollout.py
index 8c871192a..ad5dcf451 100644
--- a/tunix/rl/rollout/base_rollout.py
+++ b/tunix/rl/rollout/base_rollout.py
@@ -22,6 +22,7 @@
 from jax import numpy as jnp
 import jaxtyping
 import numpy as np
+from tunix.diffusion import types as diffusion_types
 from tunix.generate import mappings
 
 ABC = abc.ABC
@@ -63,6 +64,9 @@ class RolloutOutput:
   # The log probs from sampler generations.
   logprobs: list[np.ndarray] | None
 
+  # Optional prepared inputs for target-aligned diffusion policy scoring.
+  diffusion_batch: diffusion_types.DiffusionTokenBatch | None = None
+
 
 @dataclasses.dataclass
 class RolloutConfig:
@@ -185,7 +189,9 @@ class RolloutConfig:
   rollout_vllm_kwargs: dict[str, Any] = dataclasses.field(default_factory=dict)
 
   # Additional keyword arguments forwarded directly to the vLLM sampling params.
-  rollout_vllm_sampling_kwargs: dict[str, Any] = dataclasses.field(default_factory=dict)
+  rollout_vllm_sampling_kwargs: dict[str, Any] = dataclasses.field(
+      default_factory=dict
+  )
 
   # SG-Lang JAX specific rollout configs.
 
diff --git a/tunix/sft/diffusion.py b/tunix/sft/diffusion.py
new file mode 100644
index 000000000..f8ba2c394
--- /dev/null
+++ b/tunix/sft/diffusion.py
@@ -0,0 +1,93 @@
+# 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.
+
+"""Supervised fine-tuning loss for target-aligned diffusion scores."""
+
+from typing import Any, TypeVar
+
+from flax import nnx
+import jax.numpy as jnp
+import optax
+from tunix.diffusion import interfaces as diffusion_interfaces
+from tunix.diffusion import types as diffusion_types
+from tunix.sft import peft_trainer
+from tunix.sft import utils as sft_utils
+
+RawBatchT = TypeVar("RawBatchT")
+TrainerT = TypeVar("TrainerT", bound=peft_trainer.PeftTrainer)
+
+
+def diffusion_loss_fn(
+    model: nnx.Module,
+    batch: diffusion_types.DiffusionTokenBatch,
+    logits_fn: diffusion_interfaces.DiffusionLogitsFn,
+) -> sft_utils.LossOutput:
+  """Computes weighted cross-entropy for target-aligned diffusion scores.
+
+  ``logits_fn`` must return logits whose first two dimensions already align with
+  ``batch.target_ids``. No autoregressive one-token shift is applied. Whether
+  the batch represents CFT or SFT is determined by the external batch adapter.
+
+  Args:
+    model: The model to score.
+    batch: Canonical diffusion inputs, targets, and per-target weights.
+    logits_fn: Model-specific target-aligned logits callable.
+
+  Returns:
+    A ``LossOutput`` containing the weighted loss sum and total weight.
+  """
+
+  logits = diffusion_interfaces.compute_diffusion_logits(
+      model, batch, logits_fn
+  )
+  logits = jnp.asarray(logits, dtype=jnp.float32)
+  targets = jnp.asarray(batch.target_ids)
+  weights = jnp.asarray(batch.loss_weights, dtype=jnp.float32)
+
+  active_targets = weights != 0
+  logits = jnp.where(active_targets[..., None], logits, 0.0)
+  targets = jnp.where(active_targets, targets, 0)
+  token_losses = optax.softmax_cross_entropy_with_integer_labels(
+      logits=logits,
+      labels=targets,
+  )
+  weighted_losses = token_losses * weights
+  loss_sum = jnp.sum(weighted_losses, dtype=jnp.float32)
+  weight_sum = jnp.sum(weights, dtype=jnp.float32)
+
+  return sft_utils.LossOutput(
+      primary_loss=sft_utils.WeightedMetric(loss_sum, weight_sum),
+      aux_metrics={},
+  )
+
+
+def configure_diffusion_sft(
+    trainer: TrainerT,
+    batch_adapter: diffusion_interfaces.DiffusionBatchAdapter[RawBatchT],
+    logits_fn: diffusion_interfaces.DiffusionLogitsFn,
+) -> TrainerT:
+  """Configures and returns a trainer for diffusion supervised fine-tuning."""
+
+  def gen_model_input_fn(raw_batch: RawBatchT) -> Any:
+    return {"batch": batch_adapter(raw_batch)}
+
+  def loss_fn(
+      model: nnx.Module,
+      batch: diffusion_types.DiffusionTokenBatch,
+  ) -> sft_utils.LossOutput:
+    return diffusion_loss_fn(model, batch, logits_fn)
+
+  trainer.with_gen_model_input_fn(gen_model_input_fn)
+  trainer.with_loss_fn(loss_fn)
+  return trainer
diff --git a/tunix/sft/peft_trainer.py b/tunix/sft/peft_trainer.py
index 2c04bd4e8..a3b5afb57 100644
--- a/tunix/sft/peft_trainer.py
+++ b/tunix/sft/peft_trainer.py
@@ -63,10 +63,7 @@ class TrainingConfig:
   # contains the model params and the train data iterator state.
   checkpoint_root_directory: str | None = None
   # Checkpoint configurations. If None, the default options will be used.
-  checkpointing_options: (
-      checkpoint_options.CheckpointingOptions
-      | None
-  ) = None
+  checkpointing_options: checkpoint_options.CheckpointingOptions | None = None
 
   # Configs for the metrics logger.
   metrics_logging_options: MetricsLoggerOptions | None = None
@@ -111,6 +108,41 @@ class TrainingInput:
   images: jax.Array | np.ndarray | None = None
 
 
+def _weighted_metric_mean(values: Iterable[utils.WeightedMetric]) -> float:
+  """Aggregates unreduced metrics without microbatch-mean bias."""
+
+  values = list(values)
+  if not values:
+    return 0.0
+  if not all(isinstance(value, utils.WeightedMetric) for value in values):
+    raise TypeError("weighted metrics must not include scalar values")
+  eps = values[0].eps
+  min_denom = values[0].min_denom
+  if any(
+      value.eps != eps or value.min_denom != min_denom for value in values[1:]
+  ):
+    raise ValueError("weighted metrics must use consistent denominator bounds")
+  numerator = sum(float(np.asarray(value.unreduced_sum)) for value in values)
+  denominator = sum(float(np.asarray(value.denominator)) for value in values)
+  if eps is not None:
+    denominator += eps
+  if min_denom is not None:
+    denominator = max(denominator, min_denom)
+  return numerator / denominator if denominator else 0.0
+
+
+def _metric_reducer(
+    metric: ArrayLike | utils.WeightedMetric,
+) -> Callable[[ArrayLike], ArrayLike]:
+  """Selects the reduction that matches a buffered auxiliary metric."""
+
+  return (
+      _weighted_metric_mean
+      if isinstance(metric, utils.WeightedMetric)
+      else np.mean
+  )
+
+
 @dataclasses.dataclass(slots=True, kw_only=True)
 class MetricsBuffer:
   """Metrics collected for a specific step.
@@ -125,7 +157,7 @@ class MetricsBuffer:
   """
 
   step: int
-  losses: List[ArrayLike]
+  losses: List[ArrayLike | utils.WeightedMetric]
   additional_metrics: Dict[
       str, Tuple[List[ArrayLike], Callable[[ArrayLike], ArrayLike]]
   ] = dataclasses.field(default_factory=dict)
@@ -133,10 +165,16 @@ class MetricsBuffer:
   @property
   def loss(self):
     """Returns the mean of the recorded losses for the step."""
+    weighted = [
+        isinstance(value, utils.WeightedMetric) for value in self.losses
+    ]
+    if any(weighted):
+      if not all(weighted):
+        raise TypeError("loss values must not mix weighted and scalar metrics")
+      return _weighted_metric_mean(self.losses)
     return np.mean(np.array([np.array(x) for x in self.losses]))
 
 
-
 def _calculate_global_batch_size(train_example: Any) -> int:
   """Calculates the global batch size from a training example.
 
@@ -217,7 +255,9 @@ def _add(acc_var, g_var):
     self.denom[...] = self.denom[...] + denom_val
 
   def get(self):
-    scale = 1.0 / jnp.maximum(self.denom[...], jnp.asarray(1.0, jnp.float32))
+    denom = self.denom[...]
+    safe_denom = jnp.where(denom == 0, jnp.asarray(1.0, jnp.float32), denom)
+    scale = jnp.where(denom == 0, jnp.asarray(0.0, jnp.float32), 1 / safe_denom)
 
     return jax.tree_util.tree_map(
         lambda v: type(v)(v[...] * scale.astype(v[...].dtype)),
@@ -237,26 +277,30 @@ def _zero_in_place(v):
     self.denom[...] = jnp.zeros_like(self.denom[...])
 
 
-def _promote_opt_state_floats_to_float32(optimizer: nnx.Optimizer) -> None:
-  """Cast the optimizer state's floating-point leaves to float32 in-place.
+def _opt_state_dtypes(optimizer: nnx.Optimizer) -> Any:
+  """Returns the array dtype of every optimizer-state variable."""
+  return jax.tree_util.tree_map(
+      lambda v: v.get_value().dtype,
+      nnx.state(optimizer, nnx.optimizer.OptState),
+      is_leaf=lambda x: isinstance(x, nnx.Variable),
+  )
 
-  Args:
-    optimizer: The nnx.Optimizer instance whose state will be modified.
-  """
 
-  def _cast(v):
-    if isinstance(v, nnx.Variable):
-      val = v.value
-      if (
-          hasattr(val, "dtype")
-          and jnp.issubdtype(val.dtype, jnp.floating)
-          and val.dtype != jnp.float32
-      ):
-        v.value = val.astype(jnp.float32)
-
-  opt_state = nnx.state(optimizer, nnx.optimizer.OptState)
+def _restore_opt_state_float_dtypes(
+    optimizer: nnx.Optimizer, dtypes: Any
+) -> None:
+  """Restores floating optimizer-state leaves to their pre-update dtypes."""
+
+  def _restore(v, dtype):
+    value = v.get_value()
+    if jnp.issubdtype(value.dtype, jnp.floating) and value.dtype != dtype:
+      v.set_value(value.astype(dtype))
+
   jax.tree_util.tree_map(
-      _cast, opt_state, is_leaf=lambda x: isinstance(x, nnx.Variable)
+      _restore,
+      nnx.state(optimizer, nnx.optimizer.OptState),
+      dtypes,
+      is_leaf=lambda x: isinstance(x, nnx.Variable),
   )
 
 
@@ -309,10 +353,6 @@ def __init__(
     self._lora_enabled = utils.is_lora_enabled(self.model)
     wrt_target = nnx.LoRAParam if self._lora_enabled else nnx.Param
     self.optimizer = nnx.Optimizer(self.model, optimizer, wrt=wrt_target)
-    # Promote floating-point leaves to float32 in-place to match the dtype of
-    # the optimizer update function branch (which is float32 due to
-    # `optax.inject_hyperparams`).
-    _promote_opt_state_floats_to_float32(self.optimizer)
     self.grad_accumulator = GradientAccumulator(self.model, wrt_target)
 
     self.loss_fn = _default_loss_fn
@@ -423,7 +463,9 @@ def with_gen_model_input_fn(
       PeftTrainer.
     """
     self.clear_jit_cache()
-    self.gen_model_input_fn = gen_model_input_fn  # pyrefly: ignore[bad-assignment]
+    self.gen_model_input_fn = (
+        gen_model_input_fn  # pyrefly: ignore[bad-assignment]
+    )
     return self
 
   def _train_step(
@@ -466,16 +508,12 @@ def diff_fn(model, *args, **kwargs):
     )
     (loss_val, aux), grads = grad_fn(model, **inputs)
 
+    grad_denom = jnp.asarray(1.0, dtype=jnp.float32)
     if isinstance(aux, utils.LossOutput):
-      # Scale the unreduced gradients using the metric's scale computation
-      scale = aux.primary_loss.compute_scale()
-      grads = jax.tree.map(lambda g: g * scale, grads)
-
-      # Compute exactly equivalent legacy loss val
+      grad_denom = jnp.asarray(aux.primary_loss.denominator, dtype=jnp.float32)
       loss_val = aux.primary_loss.compute()
 
-    # TODO(b/491970038): update denom for sequence packing.
-    grad_accumulator.add(grads, denom=jnp.asarray(1.0, dtype=jnp.float32))
+    grad_accumulator.add(grads, denom=grad_denom)
 
     def apply_updates(model, optimizer, grad_accumulator):
       acc_grads = grad_accumulator.get()
@@ -486,7 +524,10 @@ def apply_updates(model, optimizer, grad_accumulator):
       norm = optax.global_norm(
           jax.tree_util.tree_map(lambda x: x.astype(jnp.float32), acc_grads)
       )
+      # `nnx.cond` requires both mutation branches to retain identical dtypes.
+      opt_state_dtypes = _opt_state_dtypes(optimizer)
       optimizer.update(model, acc_grads)
+      _restore_opt_state_float_dtypes(optimizer, opt_state_dtypes)
       grad_accumulator.reset()
       return norm
 
@@ -512,8 +553,7 @@ def skip_updates(model, optimizer, grad_accumulator):
     )
 
     if isinstance(aux, utils.LossOutput):
-      # Return the raw aux (WeightedMetric preserved); metric ops reduce them.
-      return loss_val, aux.aux_metrics, grad_norm
+      return loss_val, aux, grad_norm
     elif self._has_aux:
       return loss_val, aux, grad_norm
     else:
@@ -525,7 +565,7 @@ def _eval_step(
     inputs = self.gen_model_input_fn(inputs)
     out = self.eval_loss_fn(model, **inputs)
     if isinstance(out, utils.LossOutput):
-      return out.primary_loss.compute(), out.aux_metrics
+      return out.primary_loss.compute(), out
     elif self._has_aux:
       loss, aux = out  # pyrefly: ignore[not-iterable]
       return loss, aux
@@ -666,7 +706,9 @@ def _log_metrics(
   ):
     """Logs the metrics to the metrics logger and console."""
     perplexity = np.exp(jax.device_get(loss))
-    self.metrics_logger.log(self.metrics_prefix, "loss", loss, self._mode, step)  # pyrefly: ignore[missing-attribute]
+    self.metrics_logger.log(
+        self.metrics_prefix, "loss", loss, self._mode, step
+    )  # pyrefly: ignore[missing-attribute]
     self.metrics_logger.log(  # pyrefly: ignore[missing-attribute]
         self.metrics_prefix, "perplexity", perplexity, self._mode, step
     )
@@ -688,12 +730,14 @@ def _log_metrics(
           perplexity,
       )
     for k, v in (additional_metrics or {}).items():
-      self.metrics_logger.log(self.metrics_prefix, k, v, self._mode, step)  # pyrefly: ignore[missing-attribute]
+      self.metrics_logger.log(
+          self.metrics_prefix, k, v, self._mode, step
+      )  # pyrefly: ignore[missing-attribute]
 
   def _buffer_metrics(
       self,
       metrics_buffer: MetricsBuffer | None,
-      loss: ArrayLike,
+      loss: ArrayLike | utils.WeightedMetric,
       step: int,
       additional_metrics: (
           dict[str, Tuple[ArrayLike, Callable[[ArrayLike], ArrayLike]]] | None
@@ -707,13 +751,25 @@ def _buffer_metrics(
       )
     else:
       assert metrics_buffer.step == step
+      if isinstance(
+          metrics_buffer.losses[0], utils.WeightedMetric
+      ) != isinstance(loss, utils.WeightedMetric):
+        raise TypeError("loss values must not mix weighted and scalar metrics")
       metrics_buffer.losses.append(loss)
     if additional_metrics is not None:
       for k, (v, op) in additional_metrics.items():
         if k not in metrics_buffer.additional_metrics:
           metrics_buffer.additional_metrics[k] = ([v], op)
         else:
-          metrics_buffer.additional_metrics[k][0].append(v)
+          values = metrics_buffer.additional_metrics[k][0]
+          if isinstance(values[0], utils.WeightedMetric) != isinstance(
+              v, utils.WeightedMetric
+          ):
+            raise TypeError(
+                f"additional metric {k!r} must not mix weighted and scalar"
+                " values"
+            )
+          values.append(v)
     return metrics_buffer
 
   def _write_train_metrics(self):
@@ -744,8 +800,16 @@ def _to_np_array(v):
       return v
 
     def _apply_op(v, op):
-      if isinstance(v, list) and v and isinstance(v[0], utils.WeightedMetric):
-        if getattr(op, "__name__", "") in ("global_weighted_mean", "mean_of_means"):
+      if isinstance(v, list) and v:
+        weighted = [isinstance(x, utils.WeightedMetric) for x in v]
+        if any(weighted) and not all(weighted):
+          raise TypeError("metrics must not mix weighted and scalar values")
+      if isinstance(v, list) and v and all(weighted):
+        if getattr(op, "__name__", "") in (
+            "_weighted_metric_mean",
+            "global_weighted_mean",
+            "mean_of_means",
+        ):
           return op(v)
         v = [x.compute() for x in v]
       return op(_to_np_array(v))
@@ -937,14 +1001,27 @@ def train(
           span_v2.async_end([train_loss])
 
         self._throttler.add_computation(train_loss)
+        buffered_loss = (
+            aux.primary_loss
+            if isinstance(aux, utils.LossOutput)
+            else train_loss
+        )
+        additional_metrics = {"grad_norm": (grad_norm, np.mean)}
+        post_process_aux = aux
+        if isinstance(aux, utils.LossOutput):
+          additional_metrics.update({
+              name: (metric, _metric_reducer(metric))
+              for name, metric in aux.aux_metrics.items()
+          })
+          post_process_aux = aux.aux_metrics
         self._buffered_train_metrics = self._buffer_metrics(
             self._buffered_train_metrics,
-            loss=train_loss,
+            loss=buffered_loss,
             step=self._train_steps,
-            additional_metrics={"grad_norm": (grad_norm, np.mean)},
+            additional_metrics=additional_metrics,
         )
         # NB: put this after self._buffer_metrics is important
-        self._post_process_train_step(aux)
+        self._post_process_train_step(post_process_aux)
 
         if is_update_step_val:
           self._train_steps += 1
@@ -985,6 +1062,7 @@ def _save_last_checkpoint(self):
           self.model,
           self.optimizer,
           save_only_lora_params=self._lora_enabled,
+          custom_metadata=self.custom_checkpoint_metadata(),
           force=True,
       )
 
@@ -1026,6 +1104,8 @@ def _run_eval(
     eval_iterator = iter(eval_ds)
     with self._switch_mode(sft_metrics_logger.Mode.EVAL):
       eval_loss, eval_steps = 0, 0
+      eval_weighted_sum, eval_weight_sum = 0, 0
+      uses_weighted_loss = False
       while True:
         if self.data_hooks:
           eval_example = self.data_hooks.load_next_eval_batch(self)
@@ -1044,12 +1124,27 @@ def _run_eval(
           self.training_hooks.on_eval_step_start(self)
         loss, aux = eval_step_fn(eval_example)
         loss = jax.lax.stop_gradient(loss)
+        buffered_loss = (
+            aux.primary_loss if isinstance(aux, utils.LossOutput) else loss
+        )
+        additional_metrics = None
+        post_process_aux = aux
+        if isinstance(aux, utils.LossOutput):
+          uses_weighted_loss = True
+          eval_weighted_sum += aux.primary_loss.unreduced_sum
+          eval_weight_sum += aux.primary_loss.denominator
+          additional_metrics = {
+              name: (metric, _metric_reducer(metric))
+              for name, metric in aux.aux_metrics.items()
+          }
+          post_process_aux = aux.aux_metrics
         self._buffered_eval_metrics = self._buffer_metrics(
             self._buffered_eval_metrics,
-            loss=loss,
+            loss=buffered_loss,
             step=self._train_steps,
+            additional_metrics=additional_metrics,
         )
-        self._post_process_eval_step(aux)
+        self._post_process_eval_step(post_process_aux)
         eval_loss += loss
         eval_steps += 1
 
@@ -1059,17 +1154,23 @@ def _run_eval(
         )
         return
 
-      self._write_metrics(self._buffered_eval_metrics)  # pyrefly: ignore[bad-argument-type]
+      self._write_metrics(
+          self._buffered_eval_metrics
+      )  # pyrefly: ignore[bad-argument-type]
       logging.info(
           "Train step %d eval loss: %f - eval perplexity: %f",
           self._train_steps,
-          self.metrics_logger.get_metric(self.metrics_prefix, "loss", "eval"),  # pyrefly: ignore[missing-attribute]
+          self.metrics_logger.get_metric(
+              self.metrics_prefix, "loss", "eval"
+          ),  # pyrefly: ignore[missing-attribute]
           self.metrics_logger.get_metric(  # pyrefly: ignore[missing-attribute]
               self.metrics_prefix, "perplexity", "eval"
           ),
       )
       self._buffered_eval_metrics = None
       if self.training_hooks:
+        if uses_weighted_loss:
+          eval_loss = eval_weighted_sum / jnp.maximum(eval_weight_sum, 1)
         self.training_hooks.on_eval_step_end(self, eval_loss)
 
 
diff --git a/tunix/sft/utils.py b/tunix/sft/utils.py
index b10682513..33badedc1 100644
--- a/tunix/sft/utils.py
+++ b/tunix/sft/utils.py
@@ -229,4 +229,4 @@ class LossOutput:
   """
 
   primary_loss: WeightedMetric
-  aux_metrics: Dict[str, WeightedMetric]
+  aux_metrics: Dict[str, WeightedMetric | jax.Array]