From cc9e83c7bf0de22bbd677515c65ca0a2d737687f Mon Sep 17 00:00:00 2001
From: ethannnnnn <ethannxie@gmail.com>
Date: Wed, 22 Jul 2026 22:20:37 -0700
Subject: [PATCH 1/7] [tunix] Add framework-neutral diffusion training
 contracts

Define a target-aligned, batch-major diffusion batch contract and typed adapter/scorer protocols without depending on MaxText or a specific training algorithm.

Validate shapes and dtypes at construction and scoring boundaries, while preserving JAX pytree, JIT, and sharding compatibility.

Tests: 10 diffusion contract tests; pyink/isort; pylint; pyrefly; py_compile.
---
 tests/diffusion/types_test.py | 207 ++++++++++++++++++++++++++++++++++
 tunix/diffusion/__init__.py   |  29 +++++
 tunix/diffusion/interfaces.py |  53 +++++++++
 tunix/diffusion/types.py      | 200 ++++++++++++++++++++++++++++++++
 4 files changed, 489 insertions(+)
 create mode 100644 tests/diffusion/types_test.py
 create mode 100644 tunix/diffusion/__init__.py
 create mode 100644 tunix/diffusion/interfaces.py
 create mode 100644 tunix/diffusion/types.py

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/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

From ff2ec40dd5340c500aaf43836728c19b21f4963a Mon Sep 17 00:00:00 2001
From: ethannnnnn <ethannxie@gmail.com>
Date: Wed, 22 Jul 2026 22:56:58 -0700
Subject: [PATCH 2/7] [tunix] Accumulate weighted loss gradients globally

Accumulate LossOutput gradients as unreduced sums and normalize once by the
total denominator across microbatches. Preserve denominator-one behavior for
scalar losses and return zero gradients when every weight is zero.

Select auxiliary-metric reducers by value type in training and evaluation:
globally combine weighted metrics while averaging ordinary scalar metrics.
Reject per-key type changes across microbatches and preserve consistent
epsilon and minimum-denominator bounds during global reduction.

Preserve the dtype selected by each Optax optimizer-state initializer across
conditional update and skip branches. This keeps explicit bf16 moments in
bf16, retains explicit fp32 moments, and prevents Flax NNX branch-type
mismatches without special-casing a particular accumulation count.

Tests cover weighted and fractional denominators, zero-weight batches, mixed
weighted/plain train and eval metrics, reducer invariants, and a real
PeftTrainer + nnx.jit matrix over direct/injected AdamW and gradient
accumulation counts 1 and 2. The complete PeftTrainer suite passes 56 tests;
the cumulative focused validation passes 119 tests with six optional engine
tests deselected. Ruff and git diff checks pass.
---
 tests/sft/peft_trainer_test.py | 370 +++++++++++++++++++++++++++++----
 tunix/sft/peft_trainer.py      | 204 +++++++++++++-----
 tunix/sft/utils.py             |   2 +-
 3 files changed, 485 insertions(+), 91 deletions(-)

diff --git a/tests/sft/peft_trainer_test.py b/tests/sft/peft_trainer_test.py
index 6e4afddf4..b7b9d92d8 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):
@@ -702,10 +704,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 +823,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 +873,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 +1043,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 +1142,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 +1560,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/sft/peft_trainer.py b/tunix/sft/peft_trainer.py
index 2c04bd4e8..86bd99a59 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
@@ -1026,6 +1103,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 +1123,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 +1153,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]

From 8dc0dd39951a56720de058f2daf792eaa563e0c2 Mon Sep 17 00:00:00 2001
From: ethannnnnn <ethannxie@gmail.com>
Date: Wed, 22 Jul 2026 22:58:08 -0700
Subject: [PATCH 3/7] [tunix] Add target-aligned diffusion SFT adapter

Provide a typed PeftTrainer adapter for canonical diffusion batches and target-aligned score functions. Compute weighted float32 cross entropy without autoregressive shifting, sanitize inactive targets, and preserve zero-weight numerical safety.

Tests: 17 diffusion contract and SFT tests; 6 focused weighted-gradient tests; pyink, isort, pyrefly, pylint, py_compile, and diff checks.
---
 tests/sft/diffusion_test.py | 279 ++++++++++++++++++++++++++++++++++++
 tunix/sft/diffusion.py      |  93 ++++++++++++
 2 files changed, 372 insertions(+)
 create mode 100644 tests/sft/diffusion_test.py
 create mode 100644 tunix/sft/diffusion.py

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/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

From 6f7343e1760fc5866237a8b20139892355c37496 Mon Sep 17 00:00:00 2001
From: ethannnnnn <ethannxie@gmail.com>
Date: Wed, 22 Jul 2026 23:40:25 -0700
Subject: [PATCH 4/7] [tunix] Add prepared diffusion distillation batches

Define a framework-neutral external-teacher batch contract for freshly prepared student rollouts. Validate the canonical student batch and target-aligned teacher logits without owning model rollout, corruption, or checkpoint behavior.

Tests: 3 focused batch-contract tests; included in the 34-test diffusion contract/SFT/OPD suite.
---
 tests/distillation/diffusion_batch_test.py | 67 ++++++++++++++++++++
 tunix/distillation/diffusion.py            | 74 ++++++++++++++++++++++
 2 files changed, 141 insertions(+)
 create mode 100644 tests/distillation/diffusion_batch_test.py
 create mode 100644 tunix/distillation/diffusion.py

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/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:
+    ...

From 87538fddad45f8285bd99d0a06af3675d4726c67 Mon Sep 17 00:00:00 2001
From: ethannnnnn <ethannxie@gmail.com>
Date: Wed, 22 Jul 2026 23:41:21 -0700
Subject: [PATCH 5/7] [tunix] Add weighted prepared diffusion OPD objective

Add forward teacher-to-student KL with temperature scaling, optional target-aligned hard CE, fractional token weights, teacher stop-gradient, inactive-token sanitization, and PeftTrainer wiring for externally prepared fresh rollouts.

Tests: 11 focused OPD tests; 34 combined diffusion contract/SFT/OPD tests and 8 weighted trainer regressions passed.
---
 tests/distillation/diffusion_opd_test.py      | 390 ++++++++++++++++++
 .../train/peft_trainer_v2_test.py             |  26 +-
 tests/sft/peft_trainer_test.py                |   7 +-
 tunix/distillation/diffusion_opd.py           | 209 ++++++++++
 tunix/experimental/train/peft_trainer_v2.py   |  23 +-
 tunix/sft/peft_trainer.py                     |   1 +
 6 files changed, 638 insertions(+), 18 deletions(-)
 create mode 100644 tests/distillation/diffusion_opd_test.py
 create mode 100644 tunix/distillation/diffusion_opd.py

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/sft/peft_trainer_test.py b/tests/sft/peft_trainer_test.py
index b7b9d92d8..45479ebc0 100644
--- a/tests/sft/peft_trainer_test.py
+++ b/tests/sft/peft_trainer_test.py
@@ -571,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)
@@ -591,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
             ],
@@ -601,6 +605,7 @@ def test_checkpointing(
                 mock.ANY,
                 mock.ANY,
                 save_only_lora_params=True,
+                custom_metadata=checkpoint_metadata,
                 force=True,
             ),
             mock.call.close(),
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/sft/peft_trainer.py b/tunix/sft/peft_trainer.py
index 86bd99a59..a3b5afb57 100644
--- a/tunix/sft/peft_trainer.py
+++ b/tunix/sft/peft_trainer.py
@@ -1062,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,
       )
 

From da94a2b41ca237e9ef7097529195971337f0d3e3 Mon Sep 17 00:00:00 2001
From: ethannnnnn <ethannxie@gmail.com>
Date: Thu, 23 Jul 2026 01:14:47 -0700
Subject: [PATCH 6/7] [tunix] Add prepared diffusion policy scoring

Add a framework-neutral policy scorer for target-aligned DiffusionTokenBatch
inputs. The scorer applies the rollout temperature, sanitizes inactive padding,
optionally reports finite entropy, and makes gradient ownership explicit.
Model-aware integrations remain responsible for preparing the exact action
context or denoising trace.

Allow rollout engines to attach one optional prepared diffusion batch and merge
that pytree across rollout microbatches with strict presence, structure, and
batch-size validation. Existing autoregressive RolloutOutput construction and
generation remain unchanged when the field is absent.

Test plan:
- 16 focused diffusion/rollout generation tests passed
- 34 existing diffusion CFT/SFT/OPD tests and 6 subtests passed
- Pyink 25.12, Pylint 10/10, isort, and git diff --check
---
 tests/rl/diffusion_test.py       | 241 +++++++++++++++++++++++++++++++
 tunix/rl/diffusion.py            | 103 +++++++++++++
 tunix/rl/rl_cluster.py           |  56 +++++++
 tunix/rl/rollout/base_rollout.py |   8 +-
 4 files changed, 407 insertions(+), 1 deletion(-)
 create mode 100644 tests/rl/diffusion_test.py
 create mode 100644 tunix/rl/diffusion.py

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/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.
 

From 264e9f88b374d2d535a90de27e0bfe5bf335f77c Mon Sep 17 00:00:00 2001
From: ethannnnnn <ethannxie@gmail.com>
Date: Thu, 23 Jul 2026 02:15:43 -0700
Subject: [PATCH 7/7] [tunix] Add exact prepared diffusion policy optimization

Teach the GRPO family to optimize prepared block-diffusion rollouts without
routing them through autoregressive token scoring.

The rollout now carries its full denoising trace as a DiffusionTokenBatch.
Live, reference, and start-of-step policies score the same target-aligned
states and actions, including weighted objectives. The actor receives the
complete prepared example on its own mesh, and multi-iteration updates use a
stable old-policy snapshot.

Keep the integration opt-in through diffusion_logits_fn. Existing
autoregressive GRPO is unchanged, while PPO, agentic RL, sequence packing,
and incomplete diffusion metadata fail closed. Checkpoint metadata validates
that resumed runs retain the same diffusion objective, and custom rollout
models synchronize once at initialization without advancing global_step.

Test Plan:
- JAX_PLATFORMS=cpu XLA_FLAGS=--xla_force_host_platform_device_count=2 uv run --no-project --with-editable '.[test]' python -m pytest -q tests/rl/diffusion_test.py tests/rl/grpo/grpo_diffusion_test.py --disable-warnings
- uv run --no-project --with isort python -m isort --check-only <changed files>
- uv run --no-project --with pyink==24.10.1 python -m pyink --check <changed files>
- uv run --no-project --with-editable '.[test]' --with pylint python -m pylint --disable=all --enable=E,F <changed files>
- python3 -m py_compile <changed Python files>
- git diff --check
---
 tests/rl/diffusion_test.py             |  23 +
 tests/rl/grpo/grpo_diffusion_test.py   | 723 +++++++++++++++++++++++++
 tests/rl/rl_cluster_test.py            |  72 ++-
 tunix/rl/agentic/agentic_rl_learner.py |  68 ++-
 tunix/rl/algo_core.py                  |  72 ++-
 tunix/rl/grpo/dapo_learner.py          |   4 +
 tunix/rl/grpo/grpo_learner.py          | 299 ++++++++--
 tunix/rl/ppo/ppo_learner.py            |  18 +-
 tunix/rl/rl_cluster.py                 | 324 +++++++++--
 tunix/rl/rl_learner.py                 |  11 +-
 tunix/rl/trainer.py                    |  32 +-
 11 files changed, 1507 insertions(+), 139 deletions(-)
 create mode 100644 tests/rl/grpo/grpo_diffusion_test.py

diff --git a/tests/rl/diffusion_test.py b/tests/rl/diffusion_test.py
index 6f35f3404..7f4a9c3c4 100644
--- a/tests/rl/diffusion_test.py
+++ b/tests/rl/diffusion_test.py
@@ -195,6 +195,29 @@ def test_merges_prepared_batches_on_leading_dimension(self):
     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_preserves_host_numpy_batches_until_role_mesh_sharding(self):
+    def host_batch(target_id):
+      return diffusion_types.DiffusionTokenBatch.create(
+          model_inputs={
+              "trace_tokens": np.array([[target_id, 99]], dtype=np.int32),
+              "action_steps": np.array([[0, 1]], dtype=np.int32),
+          },
+          target_ids=np.array([[target_id, -1]], dtype=np.int32),
+          loss_weights=np.array([[1.0, 0.0]], dtype=np.float32),
+      )
+
+    merged = rl_cluster._merge_diffusion_batches(  # pylint: disable=protected-access
+        [_rollout_output(host_batch(1)), _rollout_output(host_batch(2))]
+    )
+
+    self.assertIsInstance(merged.target_ids, np.ndarray)
+    self.assertIsInstance(merged.loss_weights, np.ndarray)
+    self.assertIsInstance(merged.model_inputs["trace_tokens"], np.ndarray)
+    self.assertIsInstance(merged.model_inputs["action_steps"], np.ndarray)
+    np.testing.assert_array_equal(
+        merged.model_inputs["trace_tokens"], [[1, 99], [2, 99]]
+    )
+
   def test_allows_ar_microbatches_to_omit_diffusion_metadata(self):
     self.assertIsNone(
         rl_cluster._merge_diffusion_batches(  # pylint: disable=protected-access
diff --git a/tests/rl/grpo/grpo_diffusion_test.py b/tests/rl/grpo/grpo_diffusion_test.py
new file mode 100644
index 000000000..552476a41
--- /dev/null
+++ b/tests/rl/grpo/grpo_diffusion_test.py
@@ -0,0 +1,723 @@
+# 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 rollouts in the standard GRPO family."""
+
+# pylint: disable=missing-class-docstring,missing-function-docstring,protected-access
+
+import contextlib
+import types
+from unittest import mock
+
+from absl.testing import absltest
+from flax import nnx
+import jax
+from jax.interpreters import pxla
+import jax.numpy as jnp
+import numpy as np
+import optax
+from tunix.diffusion import types as diffusion_types
+from tunix.perf import trace as trace_lib
+from tunix.perf.experimental import tracer as perf_tracer_v2
+from tunix.rl import algo_core
+from tunix.rl import common
+from tunix.rl import diffusion
+from tunix.rl import rl_cluster as rl_cluster_lib
+from tunix.rl import rl_learner
+from tunix.rl import trainer as rl_trainer
+from tunix.rl.agentic import agentic_grpo_learner
+from tunix.rl.grpo import dapo_learner
+from tunix.rl.grpo import grpo_learner
+from tunix.rl.ppo import ppo_learner
+from tunix.rl.rollout import base_rollout
+
+
+class _LogitModel(nnx.Module):
+
+  def __init__(self, logits):
+    self.logits = nnx.Param(jnp.asarray(logits, dtype=jnp.float32))
+
+
+def _score_fn(model, model_inputs):
+  return model.logits[model_inputs["example_ids"]]
+
+
+def _batch(target_ids, loss_weights=None):
+  target_ids = jnp.asarray(target_ids, dtype=jnp.int32)
+  if loss_weights is None:
+    loss_weights = jnp.ones_like(target_ids, dtype=jnp.float32)
+  return diffusion_types.DiffusionTokenBatch.create(
+      model_inputs={
+          "example_ids": jnp.arange(target_ids.shape[0], dtype=jnp.int32),
+          "conditioning": (
+              jnp.arange(target_ids.shape[0] * 2, dtype=jnp.float32).reshape(
+                  target_ids.shape[0], 2
+              )
+          ),
+      },
+      target_ids=target_ids,
+      loss_weights=jnp.asarray(loss_weights, dtype=jnp.float32),
+  )
+
+
+def _config(*, beta=0.0, num_iterations=1):
+  config = grpo_learner.GRPOConfig(
+      num_generations=2,
+      num_iterations=num_iterations,
+      beta=beta,
+  )
+  config.temperature = 1.0
+  return config
+
+
+def _example(batch, *, old_logps=None, ref_logps=None):
+  batch_size = batch.target_ids.shape[0]
+  return grpo_learner.TrainExample(
+      prompt_ids=jnp.zeros((batch_size, 1), dtype=jnp.int32),
+      prompt_mask=jnp.ones((batch_size, 1), dtype=jnp.bool_),
+      completion_ids=batch.target_ids,
+      completion_mask=batch.loss_weights,
+      advantages=jnp.array([1.0, -0.25], dtype=jnp.float32)[:batch_size],
+      ref_per_token_logps=ref_logps,
+      old_per_token_logps=old_logps,
+      diffusion_batch=batch,
+  )
+
+
+class _FakeRollout:
+
+  def pad_id(self):
+    return 0
+
+  def eos_id(self):
+    return 3
+
+
+class _SamePadEosRollout(_FakeRollout):
+
+  def eos_id(self):
+    return 0
+
+
+class _RoutingCluster:
+
+  def __init__(self, rollout_output, ref_scores, anchor_scores):
+    self._rollout_output = rollout_output
+    self._ref_scores = ref_scores
+    self._anchor_scores = anchor_scores
+    self.ref_calls = []
+    self.anchor_calls = []
+    self.events = []
+    self.placed_examples = []
+    self.rollout = _FakeRollout()
+    self.global_steps = 0
+    self.perf = trace_lib.NoopTracer()
+    self.perf_v2 = perf_tracer_v2.NoopTracer()
+    empty_devices = types.SimpleNamespace(devices=[])
+    self.r2m = {
+        rl_cluster_lib.Role.REFERENCE: empty_devices,
+        rl_cluster_lib.Role.ACTOR: empty_devices,
+    }
+    self.cluster_config = types.SimpleNamespace(
+        rollout_config=base_rollout.RolloutConfig(
+            max_tokens_to_generate=2,
+            temperature=1.0,
+        )
+    )
+
+  def generate(self, **unused_kwargs):
+    self.events.append("generate")
+    return self._rollout_output
+
+  def snapshot_anchor_policy(self):
+    self.events.append("snapshot")
+
+  def place_pytree_on_role(self, pytree, role):
+    self.placed_examples.append((pytree, role))
+    return pytree
+
+  def get_ref_diffusion_per_token_logps(self, **kwargs):
+    self.ref_calls.append(kwargs)
+    return self._ref_scores
+
+  def get_anchor_diffusion_per_token_logps(self, **kwargs):
+    self.anchor_calls.append(kwargs)
+    return self._anchor_scores
+
+  def get_ref_per_token_logps(self, **unused_kwargs):
+    raise AssertionError("diffusion reference scoring used the AR path")
+
+  def get_old_per_token_logps(self, **unused_kwargs):
+    raise AssertionError(
+        "diffusion old-policy scoring used the rollout AR path"
+    )
+
+  def buffer_metrics(self, unused_metrics, mode):
+    del unused_metrics, mode
+
+
+class DiffusionGRPOLossTest(absltest.TestCase):
+
+  def test_live_scores_have_gradient_and_ratio_is_one_before_update(self):
+    model = _LogitModel([
+        [[2.0, -1.0, 0.0], [-1.0, 0.0, 2.0]],
+        [[0.0, 1.0, -2.0], [2.0, -1.0, 0.0]],
+    ])
+    batch = _batch([[0, 2], [1, 0]])
+    old_logps = diffusion.compute_diffusion_per_token_logps(
+        model, batch, _score_fn, stop_gradient=True
+    )
+    example = _example(batch, old_logps=old_logps)
+    config = _config()
+
+    output = algo_core.grpo_loss_fn(
+        model,
+        example,
+        config,
+        pad_id=0,
+        eos_id=3,
+        diffusion_logits_fn=_score_fn,
+    )
+    grads = nnx.grad(
+        lambda policy: algo_core.grpo_loss_fn(
+            policy,
+            example,
+            config,
+            pad_id=0,
+            eos_id=3,
+            diffusion_logits_fn=_score_fn,
+        ).primary_loss.compute()
+    )(model)
+
+    self.assertAlmostEqual(float(output.aux_metrics["is_ratio/mean"]), 1.0)
+    self.assertAlmostEqual(float(output.aux_metrics["ppo_kl"].compute()), 0.0)
+    self.assertGreater(float(jnp.linalg.norm(grads.logits[...])), 0.0)
+
+  def test_reference_kl_uses_prepared_scores(self):
+    model = _LogitModel([
+        [[2.0, 0.0], [0.0, 2.0]],
+        [[1.0, -1.0], [-1.0, 1.0]],
+    ])
+    reference = _LogitModel([
+        [[0.0, 2.0], [2.0, 0.0]],
+        [[-1.0, 1.0], [1.0, -1.0]],
+    ])
+    batch = _batch([[0, 1], [0, 1]])
+    old_logps = diffusion.compute_diffusion_per_token_logps(
+        model, batch, _score_fn
+    )
+    ref_logps = diffusion.compute_diffusion_per_token_logps(
+        reference, batch, _score_fn
+    )
+
+    output = algo_core.grpo_loss_fn(
+        model,
+        _example(batch, old_logps=old_logps, ref_logps=ref_logps),
+        _config(beta=0.1),
+        pad_id=0,
+        eos_id=3,
+        diffusion_logits_fn=_score_fn,
+    )
+
+    self.assertGreater(float(output.aux_metrics["kl"].compute()), 0.0)
+
+  def test_ar_default_still_calls_standard_scorer(self):
+    model = _LogitModel(jnp.zeros((2, 2, 3)))
+    example = grpo_learner.TrainExample(
+        prompt_ids=jnp.zeros((2, 1), dtype=jnp.int32),
+        prompt_mask=jnp.ones((2, 1), dtype=jnp.bool_),
+        completion_ids=jnp.array([[1, 2], [2, 1]], dtype=jnp.int32),
+        completion_mask=jnp.ones((2, 2), dtype=jnp.float32),
+        advantages=jnp.array([1.0, -1.0]),
+        ref_per_token_logps=None,
+        old_per_token_logps=None,
+    )
+    with mock.patch.object(
+        common,
+        "compute_per_token_logps",
+        return_value=(jnp.zeros((2, 2)), jnp.zeros((2, 2))),
+    ) as ar_scorer:
+      algo_core.grpo_loss_fn(model, example, _config(), pad_id=0, eos_id=3)
+
+    ar_scorer.assert_called_once()
+
+  def test_rejects_diffusion_batch_without_scorer(self):
+    batch = _batch([[0, 1], [1, 0]])
+    with self.assertRaisesRegex(ValueError, "refusing to score"):
+      algo_core.grpo_loss_fn(
+          _LogitModel(jnp.zeros((2, 2, 2))),
+          _example(batch),
+          _config(),
+          pad_id=0,
+          eos_id=3,
+      )
+
+
+class DiffusionGRPORoutingTest(absltest.TestCase):
+
+  def test_generation_context_hook_runs_once_before_rollout_chunks(self):
+    events = []
+
+    def generate(prompts, unused_config):
+      events.append(("generate", tuple(prompts)))
+      return base_rollout.RolloutOutput(
+          text=["x"] * len(prompts),
+          logits=None,
+          tokens=[np.array([1])] * len(prompts),
+          left_padded_prompt_tokens=np.ones((len(prompts), 1), dtype=np.int32),
+          logprobs=None,
+      )
+
+    rollout = types.SimpleNamespace(
+        model=object,
+        generate=generate,
+        set_generation_context=lambda **kwargs: events.append(
+            ("context", kwargs)
+        ),
+    )
+    cluster = object.__new__(rl_cluster_lib.RLCluster)
+    cluster._rollout = rollout
+    cluster.tokenizer = None
+    cluster.global_steps = 7
+    cluster.cluster_config = types.SimpleNamespace(
+        offload_to_cpu=False,
+        rollout_config=base_rollout.RolloutConfig(),
+    )
+    cluster._get_mesh_and_logical_axis_rules_cm = (  # pylint: disable=protected-access
+        lambda unused_role: contextlib.nullcontext(
+            (types.SimpleNamespace(devices=[]), None)
+        )
+    )
+    cluster._maybe_load_model_from_cpu = (  # pylint: disable=protected-access
+        lambda unused_model, unused_role: None
+    )
+    cluster._maybe_offload_model_to_cpu = (  # pylint: disable=protected-access
+        lambda unused_model, unused_role: None
+    )
+    cluster._perf = trace_lib.NoopTracer()
+    cluster._perf_v2 = perf_tracer_v2.NoopTracer()
+
+    cluster.generate(
+        ["p0", "p1", "p2"],
+        mode=rl_cluster_lib.Mode.EVAL,
+        micro_batch_size=2,
+    )
+
+    self.assertEqual(
+        events,
+        [
+            (
+                "context",
+                {"global_step": 7, "mode": rl_cluster_lib.Mode.EVAL},
+            ),
+            ("generate", ("p0", "p1")),
+            ("generate", ("p2",)),
+        ],
+    )
+
+  def test_learner_routes_reference_and_old_policy_through_prepared_scorer(
+      self,
+  ):
+    batch = _batch([[1, 2], [2, 1]])
+    rollout_output = base_rollout.RolloutOutput(
+        text=["a", "b"],
+        logits=None,
+        tokens=[np.array([1, 2]), np.array([2, 1])],
+        left_padded_prompt_tokens=np.array([[4], [5]]),
+        logprobs=[np.array([-9.0, -9.0]), np.array([-8.0, -8.0])],
+        diffusion_batch=batch,
+    )
+    ref_scores = jnp.full((2, 2), -0.4)
+    anchor_scores = jnp.full((2, 2), -0.5)
+    cluster = _RoutingCluster(rollout_output, ref_scores, anchor_scores)
+    learner = object.__new__(grpo_learner.GRPOLearner)
+    learner.rl_cluster = cluster
+    learner.algo_config = _config(beta=0.1, num_iterations=2)
+    learner.diffusion_logits_fn = _score_fn
+    learner.should_sync_weights = False
+    learner._rollout_micro_batch_size = 1
+    learner._compute_logps_micro_batch_size = 1
+    learner.metric_fns = []
+    learner._compute_rewards = lambda **unused_kwargs: np.array([1.0, 2.0])
+
+    result = learner._generate_and_compute_advantage({"prompts": ["p0", "p1"]})
+
+    self.assertLen(cluster.ref_calls, 1)
+    self.assertLen(cluster.anchor_calls, 1)
+    self.assertEqual(cluster.events[:2], ["snapshot", "generate"])
+    self.assertIs(result.diffusion_batch, batch)
+    self.assertLen(cluster.placed_examples, 1)
+    self.assertIs(cluster.placed_examples[0][1], rl_cluster_lib.Role.ACTOR)
+    np.testing.assert_array_equal(result.ref_per_token_logps, ref_scores)
+    np.testing.assert_array_equal(result.old_per_token_logps, anchor_scores)
+
+  def test_prepared_targets_and_weights_must_match_completion(self):
+    with self.assertRaisesRegex(ValueError, "target_ids must exactly match"):
+      grpo_learner._validate_diffusion_rollout_batch(  # pylint: disable=protected-access
+          _batch([[1, 0], [2, 0]], [[1, 0], [1, 0]]),
+          [np.array([2]), np.array([2])],
+      )
+    with self.assertRaisesRegex(ValueError, "active loss_weights"):
+      grpo_learner._validate_diffusion_rollout_batch(  # pylint: disable=protected-access
+          _batch([[1, 0], [2, 0]], [[0, 1], [1, 0]]),
+          [np.array([1]), np.array([2])],
+      )
+
+  def test_inactive_targets_can_retain_full_trace_sentinels(self):
+    batch = _batch([[1, -1], [2, 99]], [[1, 0], [1, 0]])
+
+    result = grpo_learner._validate_diffusion_rollout_batch(  # pylint: disable=protected-access
+        batch,
+        [np.array([1]), np.array([2])],
+    )
+
+    self.assertIs(result, batch)
+
+  def test_sampler_parity_requires_full_trace_logprobs(self):
+    with self.assertRaisesRegex(ValueError, "full prepared trace"):
+      grpo_learner._stack_rollout_logps(  # pylint: disable=protected-access
+          [np.array([-0.1, -0.2])], batch_size=1, trace_length=3
+      )
+
+  def test_missing_prepared_batch_fails_instead_of_falling_back(self):
+    rollout_output = base_rollout.RolloutOutput(
+        text=["a", "b"],
+        logits=None,
+        tokens=[np.array([1]), np.array([2])],
+        left_padded_prompt_tokens=np.array([[4], [5]]),
+        logprobs=None,
+    )
+    cluster = _RoutingCluster(
+        rollout_output, jnp.zeros((2, 2)), jnp.zeros((2, 2))
+    )
+    learner = object.__new__(grpo_learner.GRPOLearner)
+    learner.rl_cluster = cluster
+    learner.algo_config = _config()
+    learner.diffusion_logits_fn = _score_fn
+    learner.should_sync_weights = True
+    learner._rollout_micro_batch_size = 1
+    learner._compute_logps_micro_batch_size = 1
+    learner.metric_fns = []
+
+    with self.assertRaisesRegex(ValueError, "requires every rollout"):
+      learner._generate_and_compute_advantage({"prompts": ["p0", "p1"]})
+
+  def test_full_trace_survives_eos_when_pad_and_eos_ids_match(self):
+    batch = _batch([[1, 0, 7], [2, 0, 8]])
+    rollout_output = base_rollout.RolloutOutput(
+        text=["a", "b"],
+        logits=None,
+        tokens=[np.array([1, 0]), np.array([2, 0])],
+        left_padded_prompt_tokens=np.array([[4], [5]]),
+        logprobs=None,
+        diffusion_batch=batch,
+    )
+    cluster = _RoutingCluster(
+        rollout_output, jnp.zeros((2, 3)), jnp.zeros((2, 3))
+    )
+    cluster.rollout = _SamePadEosRollout()
+    learner = object.__new__(grpo_learner.GRPOLearner)
+    learner.rl_cluster = cluster
+    learner.algo_config = _config()
+    learner.diffusion_logits_fn = _score_fn
+    learner.should_sync_weights = True
+    learner._rollout_micro_batch_size = 1
+    learner._compute_logps_micro_batch_size = 1
+    learner.metric_fns = []
+    learner._compute_rewards = lambda **unused_kwargs: np.array([1.0, 2.0])
+
+    result = learner._generate_and_compute_advantage({"prompts": ["p0", "p1"]})
+
+    np.testing.assert_array_equal(result.completion_ids, [[1, 0, 7], [2, 0, 8]])
+    np.testing.assert_array_equal(result.completion_mask, np.ones((2, 3)))
+
+  def test_ppo_rejects_prepared_diffusion_rollout(self):
+    rollout_output = base_rollout.RolloutOutput(
+        text=["a"],
+        logits=None,
+        tokens=[np.array([1])],
+        left_padded_prompt_tokens=np.array([[4]]),
+        logprobs=None,
+        diffusion_batch=_batch([[1]]),
+    )
+    learner = object.__new__(ppo_learner.PPOLearner)
+    learner.rl_cluster = types.SimpleNamespace(
+        cluster_config=types.SimpleNamespace(
+            rollout_config=base_rollout.RolloutConfig()
+        ),
+        rollout=_FakeRollout(),
+        generate=lambda **unused_kwargs: rollout_output,
+    )
+    learner._rollout_micro_batch_size = 1
+
+    with self.assertRaisesRegex(ValueError, "not supported by PPO"):
+      learner._generate_and_compute_advantage({"prompts": ["p0"]})
+
+  def test_agentic_rl_rejects_prepared_diffusion_rollout(self):
+    rollout_output = base_rollout.RolloutOutput(
+        text=["a"],
+        logits=None,
+        tokens=[np.array([1])],
+        left_padded_prompt_tokens=np.array([[4]]),
+        logprobs=None,
+        diffusion_batch=_batch([[1]]),
+    )
+    learner = object.__new__(agentic_grpo_learner.GRPOLearner)
+    learner.rl_cluster = types.SimpleNamespace(
+        generate=lambda **unused_kwargs: rollout_output
+    )
+    learner.chat_parser = None
+    learner.policy_version = 0
+    learner._full_batch_size = 0
+
+    with self.assertRaisesRegex(ValueError, "not supported by agentic RL"):
+      learner._model_call([])
+
+
+class DiffusionPolicySnapshotTest(absltest.TestCase):
+
+  def _bare_cluster(self):
+    cluster = object.__new__(rl_cluster_lib.RLCluster)
+    cluster.cluster_config = types.SimpleNamespace(
+        offload_to_cpu=False,
+        training_config=types.SimpleNamespace(data_sharding_axis=()),
+    )
+    cluster.r2m = {
+        rl_cluster_lib.Role.ACTOR: pxla.thread_resources.env.physical_mesh,
+        rl_cluster_lib.Role.REFERENCE: pxla.thread_resources.env.physical_mesh,
+    }
+    cluster._get_mesh_and_logical_axis_rules_cm = (  # pylint: disable=protected-access
+        lambda unused_role: contextlib.nullcontext()
+    )
+    cluster._maybe_load_model_from_cpu = (  # pylint: disable=protected-access
+        lambda unused_model, unused_role: None
+    )
+    cluster._maybe_offload_model_to_cpu = (  # pylint: disable=protected-access
+        lambda unused_model, unused_role: None
+    )
+    cluster.get_rollout_config = lambda mode: base_rollout.RolloutConfig(
+        temperature=1.0
+    )
+    return cluster
+
+  def test_reference_scorer_microbatches_the_whole_prepared_pytree(self):
+    model = _LogitModel([
+        [[2.0, 0.0]],
+        [[0.0, 2.0]],
+        [[1.0, -1.0]],
+    ])
+    batch = _batch([[0], [1], [0]])
+    seen_shapes = []
+
+    def score_fn(policy, model_inputs):
+      seen_shapes.append(
+          tuple(leaf.shape[0] for leaf in jax.tree.leaves(model_inputs))
+      )
+      return _score_fn(policy, model_inputs)
+
+    cluster = self._bare_cluster()
+    cluster._inference_worker = types.SimpleNamespace(  # pylint: disable=protected-access
+        get_model=lambda unused_name: model
+    )
+    scores = cluster.get_ref_diffusion_per_token_logps(
+        batch, score_fn, micro_batch_size=2
+    )
+
+    self.assertEqual(seen_shapes, [(2, 2), (1, 1)])
+    expected = diffusion.compute_diffusion_per_token_logps(
+        model, batch, _score_fn
+    )
+    np.testing.assert_allclose(scores, expected)
+
+  def test_anchor_scorer_uses_start_of_step_state_after_live_model_changes(
+      self,
+  ):
+    initial_logits = jnp.array([
+        [[2.0, 0.0]],
+        [[0.0, 2.0]],
+        [[1.0, -1.0]],
+    ])
+    live_model = _LogitModel(initial_logits)
+    _, initial_state = nnx.split(live_model)
+    anchor_state = jax.tree.map(jnp.copy, initial_state)
+    live_model.logits[...] = -initial_logits
+    batch = _batch([[0], [1], [0]])
+
+    cluster = self._bare_cluster()
+    cluster._actor_trainer = types.SimpleNamespace(  # pylint: disable=protected-access
+        model=live_model
+    )
+    cluster._anchor_policy_state = anchor_state  # pylint: disable=protected-access
+    cluster._default_memory_kind = jax.devices()[0].default_memory().kind  # pylint: disable=protected-access
+    cluster._is_state_on_device = lambda unused_state: True  # pylint: disable=protected-access
+    scores = cluster.get_anchor_diffusion_per_token_logps(
+        batch, _score_fn, micro_batch_size=2
+    )
+
+    expected = diffusion.compute_diffusion_per_token_logps(
+        _LogitModel(initial_logits), batch, _score_fn
+    )
+    live_scores = diffusion.compute_diffusion_per_token_logps(
+        live_model, batch, _score_fn
+    )
+    np.testing.assert_allclose(scores, expected)
+    self.assertFalse(np.allclose(scores, live_scores))
+
+  def test_dapo_forwards_diffusion_scorer(self):
+    scorer = mock.Mock()
+    with mock.patch.object(grpo_learner.GRPOLearner, "__init__") as base_init:
+      dapo_learner.DAPOLearner(
+          rl_cluster=mock.Mock(),
+          algo_config=dapo_learner.DAPOConfig(overlong_buffer=None),
+          reward_fns=lambda **unused_kwargs: [0.0],
+          diffusion_logits_fn=scorer,
+      )
+
+    self.assertIs(base_init.call_args.kwargs["diffusion_logits_fn"], scorer)
+
+  def test_complete_train_example_is_placed_on_actor_mesh(self):
+    if jax.device_count() < 2:
+      self.skipTest("requires two devices to construct disjoint role meshes")
+    actor_mesh = jax.sharding.Mesh(np.asarray([jax.devices()[0]]), ("fsdp",))
+    other_mesh = jax.sharding.Mesh(np.asarray([jax.devices()[1]]), ("fsdp",))
+    other_sharding = jax.sharding.NamedSharding(
+        other_mesh, jax.sharding.PartitionSpec()
+    )
+    batch = jax.tree.map(
+        lambda value: jax.device_put(value, other_sharding),
+        _batch([[0, 1], [1, 0]]),
+    )
+    ref_logps = jax.device_put(jnp.full((2, 2), -0.5), other_sharding)
+    example = _example(batch, ref_logps=ref_logps)
+
+    cluster = object.__new__(rl_cluster_lib.RLCluster)
+    cluster.r2m = {rl_cluster_lib.Role.ACTOR: actor_mesh}
+    cluster.cluster_config = types.SimpleNamespace(
+        training_config=types.SimpleNamespace(data_sharding_axis=("fsdp",))
+    )
+    cluster._get_mesh_and_logical_axis_rules_cm = (  # pylint: disable=protected-access
+        lambda role: cluster.r2m[role]
+    )
+    placed = cluster.place_pytree_on_role(example, rl_cluster_lib.Role.ACTOR)
+
+    for leaf in jax.tree.leaves(placed):
+      self.assertIsInstance(leaf, jax.Array)
+      self.assertEqual(leaf.sharding.mesh, actor_mesh)
+
+  def test_disjoint_rollout_is_initially_synced_without_advancing_step(self):
+    class MinimalLearner(rl_learner.RLLearner):
+
+      def _generate_and_compute_advantage(self, training_input, mode):
+        del training_input, mode
+
+      def _compute_trajectory_ids(self, example, steps):
+        del example, steps
+        return []
+
+      def _num_iterations(self):
+        return 1
+
+      def _num_generations(self):
+        return 1
+
+    actor_trainer = types.SimpleNamespace(
+        model=_LogitModel(jnp.zeros((1, 1, 2))),
+        restored_global_step=lambda: 7,
+        iter_steps=0,
+        is_managed_externally=False,
+    )
+    cluster = types.SimpleNamespace(
+        actor_trainer=actor_trainer,
+        rollout=types.SimpleNamespace(
+            model=lambda: _LogitModel(jnp.ones((1, 1, 2)))
+        ),
+        cluster_config=types.SimpleNamespace(
+            training_config=types.SimpleNamespace(
+                rollout_micro_batch_size=None,
+                compute_logps_micro_batch_size=None,
+            ),
+            role_to_mesh={
+                rl_cluster_lib.Role.ACTOR: "actor",
+                rl_cluster_lib.Role.ROLLOUT: "rollout",
+            },
+        ),
+        global_steps=0,
+        sync_weights=mock.Mock(),
+    )
+    with mock.patch(
+        "tunix.rl.rl_learner.rl_utils.is_sharing_weights", return_value=False
+    ), mock.patch("tunix.rl.rl_learner.sft_utils.show_hbm_usage"):
+      learner = MinimalLearner(
+          rl_cluster=cluster,
+          algo_config=_config(),
+          reward_fns=lambda **unused_kwargs: [0.0],
+      )
+
+    cluster.sync_weights.assert_called_once_with(increment_global_steps=False)
+    self.assertEqual(cluster.global_steps, 7)
+    learner.executor.shutdown()
+
+
+class RLCheckpointMetadataTest(absltest.TestCase):
+
+  def test_empty_metadata_default_preserves_existing_behavior(self):
+    config = rl_cluster_lib.RLTrainingConfig(
+        actor_optimizer=optax.sgd(1e-3),
+        eval_every_n_steps=1,
+    )
+    self.assertEmpty(config.checkpoint_metadata)
+
+    trainer = object.__new__(rl_trainer.Trainer)
+    trainer._train_steps = 4
+    trainer._restored_custom_metadata = {"global_step": 4}
+    trainer._validate_restored_checkpoint_metadata({})  # pylint: disable=protected-access
+
+  def test_restored_metadata_requires_every_matching_value(self):
+    trainer = object.__new__(rl_trainer.Trainer)
+    trainer._train_steps = 4
+    trainer._restored_custom_metadata = {
+        "global_step": 4,
+        "objective": "diffusion-grpo",
+        "stop_ids": [1, 2],
+    }
+    trainer._validate_restored_checkpoint_metadata(  # pylint: disable=protected-access
+        {"objective": "diffusion-grpo", "stop_ids": (1, 2)}
+    )
+
+    with self.assertRaisesRegex(ValueError, "mismatch"):
+      trainer._validate_restored_checkpoint_metadata(  # pylint: disable=protected-access
+          {"objective": "autoregressive"}
+      )
+    with self.assertRaisesRegex(ValueError, "missing required"):
+      trainer._validate_restored_checkpoint_metadata(  # pylint: disable=protected-access
+          {"mask_id": 123}
+      )
+
+  def test_metadata_config_rejects_reserved_and_non_json_values(self):
+    with self.assertRaisesRegex(ValueError, "reserved"):
+      rl_cluster_lib.RLTrainingConfig(
+          actor_optimizer=optax.sgd(1e-3),
+          eval_every_n_steps=1,
+          checkpoint_metadata={"role": "actor"},
+      )
+    with self.assertRaisesRegex(TypeError, "JSON scalar or tuple"):
+      rl_cluster_lib.RLTrainingConfig(
+          actor_optimizer=optax.sgd(1e-3),
+          eval_every_n_steps=1,
+          checkpoint_metadata={"bad": [1, 2]},
+      )
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/rl/rl_cluster_test.py b/tests/rl/rl_cluster_test.py
index 13b53f9c5..014cb63c3 100644
--- a/tests/rl/rl_cluster_test.py
+++ b/tests/rl/rl_cluster_test.py
@@ -14,7 +14,6 @@
 
 import functools
 import os
-import os
 import unittest
 from unittest import mock
 
@@ -134,6 +133,71 @@ def test_model_loading_with_resharding(self):
     )
     self.assertEqual(ref_model_mesh, actor_mesh)
 
+  def test_model_loading_maps_logical_axes_to_role_mesh(self):
+    split_index = self.device_count // 2
+    actor_mesh = Mesh(
+        np.array(jax.devices()[:split_index]).reshape(split_index, 1),
+        ('data', 'model'),
+    )
+    rollout_mesh = Mesh(
+        np.array(jax.devices()[split_index:]).reshape(1, split_index),
+        ('data', 'model'),
+    )
+    logical_axis_rules = (('fsdp', 'data'), ('tp', 'model'))
+    cluster_config = rl_cluster_lib.ClusterConfig(
+        role_to_mesh={
+            rl_cluster_lib.Role.ACTOR: actor_mesh,
+            rl_cluster_lib.Role.REFERENCE: actor_mesh,
+            rl_cluster_lib.Role.ROLLOUT: rollout_mesh,
+        },
+        role_to_logical_axis_rule={
+            rl_cluster_lib.Role.ACTOR: logical_axis_rules,
+            rl_cluster_lib.Role.REFERENCE: logical_axis_rules,
+            rl_cluster_lib.Role.ROLLOUT: logical_axis_rules,
+        },
+        rollout_engine='vanilla',
+        offload_to_cpu=False,
+        training_config=rl_cluster_lib.RLTrainingConfig(
+            actor_optimizer=optax.sgd(1e-3),
+            eval_every_n_steps=1,
+            max_steps=10,
+            gradient_accumulation_steps=None,
+        ),
+        rollout_config=base_rollout.RolloutConfig(
+            max_tokens_to_generate=10,
+            max_prompt_length=256,
+            kv_cache_size=1024,
+            data_type=jnp.bfloat16,
+        ),
+    )
+    vocab = tc.MockVocab()
+    model = tc.ToyTransformer(
+        config=tc.ModelConfig(vocab_size=vocab.GetPieceSize()), rngs=nnx.Rngs(0)
+    )
+    ref_model = tc.ToyTransformer(
+        config=tc.ModelConfig(vocab_size=vocab.GetPieceSize()), rngs=nnx.Rngs(0)
+    )
+
+    cluster = rl_cluster_lib.RLCluster(
+        actor=model,
+        reference=ref_model,
+        tokenizer=vocab,
+        cluster_config=cluster_config,
+    )
+
+    states = (
+        nnx.state(cluster.actor_trainer.model),
+        nnx.state(cluster.rollout.model()),
+        nnx.state(cluster.inference_worker._models['reference']),
+    )
+    for state in states:
+      for sharding in jax.tree.leaves(
+          jax.tree.map(lambda x: x.sharding, state)
+      ):
+        self.assertNotIn('fsdp', sharding.spec)
+        self.assertNotIn('tp', sharding.spec)
+        self.assertTrue(set(sharding.spec) <= {'data', 'model', None})
+
   @parameterized.named_parameters(
       dict(
           testcase_name='2d_mesh_perf_v1_only',
@@ -788,16 +852,16 @@ def create_cluster_config(rollout_engine):
       dict(
           testcase_name='missing_role',
           role_to_logical_axis_rules={
-              rl_cluster_lib.Role.ACTOR: ['fsdp'],
+              rl_cluster_lib.Role.ACTOR: [('fsdp', 'fsdp')],
           },
           expected_logical_axis_rules=(),
       ),
       dict(
           testcase_name='with_rule',
           role_to_logical_axis_rules={
-              rl_cluster_lib.Role.REFERENCE: ['fsdp'],
+              rl_cluster_lib.Role.REFERENCE: [('fsdp', 'fsdp')],
           },
-          expected_logical_axis_rules=['fsdp'],
+          expected_logical_axis_rules=[('fsdp', 'fsdp')],
       ),
   )
   def test_logical_axis_rules_cm(
diff --git a/tunix/rl/agentic/agentic_rl_learner.py b/tunix/rl/agentic/agentic_rl_learner.py
index b698ee57e..a6aa32142 100644
--- a/tunix/rl/agentic/agentic_rl_learner.py
+++ b/tunix/rl/agentic/agentic_rl_learner.py
@@ -15,8 +15,8 @@
 """Base class for Agentic RL Learners."""
 
 from __future__ import annotations
+
 import abc
-import time
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
 import contextlib
@@ -25,7 +25,8 @@
 import itertools
 import queue
 import threading
-from typing import Any, AsyncIterator, Callable, Dict, Generic, Iterable, Iterator, List, Sequence, Type, TypeVar, Optional, Set
+import time
+from typing import Any, AsyncIterator, Callable, Dict, Generic, Iterable, Iterator, List, Optional, Sequence, Set, Type, TypeVar
 
 from absl import logging
 import flax
@@ -33,13 +34,12 @@
 from jax import typing
 import jax.numpy as jnp
 import numpy as np
+from tunix.perf.experimental import constants as perf_constants
 from tunix.rl import algorithm_config as algo_config_lib
 from tunix.rl import common
-from tunix.perf.experimental import constants as perf_constants
 from tunix.rl import function_registry
 from tunix.rl import reward_manager  # pylint: disable=unused-import
 from tunix.rl import rl_cluster as rl_cluster_lib
-from tunix.rl.rollout import base_rollout
 from tunix.rl import utils as rl_utils
 from tunix.rl.agentic import utils as agentic_utils
 from tunix.rl.agentic.agents import base_agent
@@ -50,6 +50,7 @@
 from tunix.rl.agentic.rewards import reward  # pylint: disable=unused-import
 from tunix.rl.agentic.trajectory import trajectory_collect_engine
 from tunix.rl.queue import data_queue as queue_lib
+from tunix.rl.rollout import base_rollout
 from tunix.sft import utils as sft_utils
 
 ArrayLike = typing.ArrayLike
@@ -422,7 +423,11 @@ def _create_agent_env_pair(
     assert "pair_index" not in self.env_kwargs
     env = self.env_class(
         single_example,
-        **{"group_id": group_id, "pair_index": pair_index, **self.env_kwargs},  # pyrefly: ignore[bad-argument-type]
+        **{
+            "group_id": group_id,
+            "pair_index": pair_index,
+            **self.env_kwargs,
+        },  # pyrefly: ignore[bad-argument-type]
     )
 
     return agent, env
@@ -461,6 +466,11 @@ def _model_call(
         trace_tags=tags,
         max_generation_steps=max_generation_steps,
     )
+    if result.diffusion_batch is not None:
+      raise ValueError(
+          "prepared diffusion rollouts are not supported by agentic RL; use "
+          "the standard GRPO family with a diffusion_logits_fn"
+      )
 
     return result
 
@@ -509,7 +519,9 @@ async def pairs_stream_generator():
       # with mini-batch.
       group_id = self.rl_cluster.global_steps * self._full_batch_size
       if is_async_iterator:
-        async for single_example in prompt_iterator:  # pyrefly: ignore[not-iterable]
+        async for (
+            single_example
+        ) in prompt_iterator:  # pyrefly: ignore[not-iterable]
           # Create agent-env pairs in parallel for a group to handle potential
           # cold start latency on env creation.
           agent_env_pairs = await asyncio.gather(*[
@@ -726,7 +738,9 @@ def train(
       self.rl_cluster.close()
       return
 
-    full_batch_size = len(next(iter(first_item.values())))  # pyrefly: ignore[bad-argument-type]
+    full_batch_size = len(
+        next(iter(first_item.values()))
+    )  # pyrefly: ignore[bad-argument-type]
     self._full_batch_size = full_batch_size
     # Initialize batch sizes.
     mini_batch_size = self._training_config.mini_batch_size or full_batch_size
@@ -865,9 +879,7 @@ def train(
       # example along its batch axis into chunks sized to one micro-step,
       # and pass the list to ``update_actor``; ``peft_trainer.train``
       # iterates the list and calls ``train_step`` once per chunk.
-      seqs_per_chunk = (
-          train_micro_batch_size * self.algo_config.num_generations
-      )
+      seqs_per_chunk = train_micro_batch_size * self.algo_config.num_generations
       n_total = merged_train_micro_batch.completion_ids.shape[0]
       if n_total > seqs_per_chunk:
         chunked_train_micro_batch = [
@@ -946,10 +958,10 @@ async def _eval_runner_async(current_eval_orchestrator):
         # `is_update_step` flips True every `grad_acc_steps` micro-batches.
         unpacked_micro_step_counter += 1
         is_update = unpacked_micro_step_counter % grad_acc_steps == 0
-        
+
       if is_update:
         update_steps_since_last_sync += 1
-        
+
       if update_steps_since_last_sync == update_steps_per_full_batch:
         # --- Remaining Iterations Training Step ---
         iterations = self._num_iterations()
@@ -958,18 +970,13 @@ async def _eval_runner_async(current_eval_orchestrator):
           # TODO(b/483779605) Sub-step checkpointing.
           self._iter_steps += len(full_batch_chunks)
 
-          # TODO(yixuanm): Eval during iteration too. Skipping for now as we 
+          # TODO(yixuanm): Eval during iteration too. Skipping for now as we
           # will refactor the learner soon.
-          self.rl_cluster.update_actor(
-              full_batch_chunks, None, skip_jit
-          )
+          self.rl_cluster.update_actor(full_batch_chunks, None, skip_jit)
           if hasattr(self.rl_cluster, "critic_trainer"):
-            self.rl_cluster.update_critic(
-                full_batch_chunks, None, skip_jit
-            )
+            self.rl_cluster.update_critic(full_batch_chunks, None, skip_jit)
         full_batch_chunks.clear()
 
-
         global_step_time = time.time() - self._global_step_start_time
         logging.info(
             f"Global step {self.rl_cluster.global_steps} completed in"
@@ -980,7 +987,9 @@ async def _eval_runner_async(current_eval_orchestrator):
         # step). Mirrors the per-iter view a wandb dashboard would show
         # without depending on the async metric logger pipeline.
         with self._rewards_window_lock:
-          train_rewards = np.asarray(self._train_rewards_window, dtype=np.float32)
+          train_rewards = np.asarray(
+              self._train_rewards_window, dtype=np.float32
+          )
           eval_rewards = np.asarray(self._eval_rewards_window, dtype=np.float32)
           self._train_rewards_window.clear()
           if did_eval_this_global_step:
@@ -1016,10 +1025,9 @@ async def _eval_runner_async(current_eval_orchestrator):
         trainer_str = ""
         try:
           actor_trainer = self.rl_cluster.actor_trainer
-          trainer_buf = (
-              getattr(actor_trainer, "_prev_buffered_train_metrics", None)
-              or getattr(actor_trainer, "_buffered_train_metrics", None)
-          )
+          trainer_buf = getattr(
+              actor_trainer, "_prev_buffered_train_metrics", None
+          ) or getattr(actor_trainer, "_buffered_train_metrics", None)
           if trainer_buf is not None:
             extras = []
             if trainer_buf.losses:
@@ -1037,9 +1045,9 @@ async def _eval_runner_async(current_eval_orchestrator):
                 vals, _ = am[key]
                 if vals:
                   v = float(
-                      np.mean([
-                          np.asarray(common._metric_scalar(x)) for x in vals
-                      ])
+                      np.mean(
+                          [np.asarray(common._metric_scalar(x)) for x in vals]
+                      )
                   )
                   extras.append(f"{label}={v:.4f}")
             if extras:
@@ -1182,7 +1190,9 @@ def _filter_outdated_offpolicy_examples(
           len(train_micro_batch),
           self.policy_version,
           str([
-              train_example.policy_version[0]  # pyrefly: ignore[unsupported-operation]
+              train_example.policy_version[
+                  0
+              ]  # pyrefly: ignore[unsupported-operation]
               for train_example in train_micro_batch
           ]),
           self.algo_config.off_policy_steps,
diff --git a/tunix/rl/algo_core.py b/tunix/rl/algo_core.py
index ad2f234fa..96f994d13 100644
--- a/tunix/rl/algo_core.py
+++ b/tunix/rl/algo_core.py
@@ -15,11 +15,14 @@
 """Algorithm core implementations for RL and Agentic RL learners."""
 
 import functools
+
 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 common
+from tunix.rl import diffusion as diffusion_lib
 from tunix.rl import function_registry
 from tunix.sft import utils as sft_utils
 
@@ -200,7 +203,6 @@ def ppo_policy_loss_fn(
   else:
     per_token_logps = outputs
 
-
   advantages = train_example.advantages
   old_per_token_logps = train_example.old_per_token_logps
 
@@ -404,22 +406,50 @@ def grpo_loss_fn(
       train_example.completion_mask,
   )
 
-  # TODO(tsbao): split can be avoided with updated peft_trainer model handling.
-  graphdef, state = nnx.split(model)
-  per_token_logps, token_entropy = common.compute_per_token_logps(
-      graphdef,
-      state,
-      prompt_tokens=train_example.prompt_ids,
-      completion_tokens=completion_ids,
-      pad_id=pad_id,
-      eos_id=eos_id,
-      stop_gradient=False,
-      return_entropy=True,
-      segment_ids=getattr(train_example, "segment_ids", None),
-      segment_positions=getattr(train_example, "segment_positions", None),
-      temperature=algo_config.temperature,
-      chunk_size=kwargs.get("compute_logps_chunk_size", 0),
-  )
+  diffusion_logits_fn = kwargs.get("diffusion_logits_fn")
+  diffusion_batch = getattr(train_example, "diffusion_batch", None)
+  if diffusion_logits_fn is None:
+    if isinstance(diffusion_batch, diffusion_types.DiffusionTokenBatch):
+      raise ValueError(
+          "diffusion_batch requires a diffusion_logits_fn; refusing to score "
+          "a diffusion rollout autoregressively"
+      )
+    # TODO(tsbao): split can be avoided with updated peft_trainer model handling.
+    graphdef, state = nnx.split(model)
+    per_token_logps, token_entropy = common.compute_per_token_logps(
+        graphdef,
+        state,
+        prompt_tokens=train_example.prompt_ids,
+        completion_tokens=completion_ids,
+        pad_id=pad_id,
+        eos_id=eos_id,
+        stop_gradient=False,
+        return_entropy=True,
+        segment_ids=getattr(train_example, "segment_ids", None),
+        segment_positions=getattr(train_example, "segment_positions", None),
+        temperature=algo_config.temperature,
+        chunk_size=kwargs.get("compute_logps_chunk_size", 0),
+    )
+  else:
+    if not isinstance(diffusion_batch, diffusion_types.DiffusionTokenBatch):
+      raise ValueError(
+          "diffusion_logits_fn requires TrainExample.diffusion_batch"
+      )
+    if (
+        getattr(train_example, "segment_ids", None) is not None
+        or getattr(train_example, "segment_positions", None) is not None
+    ):
+      raise ValueError("diffusion GRPO does not support sequence packing")
+    per_token_logps, token_entropy = (
+        diffusion_lib.compute_diffusion_per_token_logps(
+            model,
+            diffusion_batch,
+            diffusion_logits_fn,
+            temperature=algo_config.temperature,
+            stop_gradient=False,
+            return_entropy=True,
+        )
+    )
   per_token_logps = jnp.astype(per_token_logps, jnp.float32)
   # TODO(tsbao): We should handle token level advantages.
   advantages = jnp.astype(train_example.advantages, jnp.float32)
@@ -503,13 +533,13 @@ def grpo_loss_fn(
   reduced_pg_loss = common.reduced_loss_agg(
       per_token_loss, completion_mask, loss_aggregation_mode
   )
-  total_loss = unreduced_pg_loss  # KL added below when beta != 0; feeds gradient
+  total_loss = (
+      unreduced_pg_loss  # KL added below when beta != 0; feeds gradient
+  )
   # Per-token diagnostics — log only over assistant tokens (completion_mask).
   is_ratio_mean = masked_mean(is_ratio, completion_mask)
   is_ratio_max = jnp.max(jnp.where(completion_mask > 0, is_ratio, 0.0))
-  is_ratio_min = jnp.min(
-      jnp.where(completion_mask > 0, is_ratio, jnp.inf)
-  )
+  is_ratio_min = jnp.min(jnp.where(completion_mask > 0, is_ratio, jnp.inf))
   log_ratio_abs_mean = masked_mean(
       jnp.abs(seq_importance_ratio), completion_mask
   )
diff --git a/tunix/rl/grpo/dapo_learner.py b/tunix/rl/grpo/dapo_learner.py
index 6e48d2e5c..8164eeb00 100644
--- a/tunix/rl/grpo/dapo_learner.py
+++ b/tunix/rl/grpo/dapo_learner.py
@@ -15,6 +15,8 @@
 
 import dataclasses
 from typing import Any, Dict, List, Optional, Sequence
+
+from tunix.diffusion import interfaces as diffusion_interfaces
 from tunix.rl import rl_cluster as rl_cluster_lib
 from tunix.rl import rl_learner
 from tunix.rl.grpo import grpo_learner as grpo_learner_lib
@@ -140,6 +142,7 @@ def __init__(
       reward_fns: RewardFn | List[RewardFn],
       metric_fns: Sequence[MetricFn] | None = None,
       data_shuffle_seed: int | None = None,
+      diffusion_logits_fn: diffusion_interfaces.DiffusionLogitsFn | None = None,
   ):
     """Initializes the `DAPOLearner`."""
     reward_fns = (
@@ -153,4 +156,5 @@ def __init__(
         reward_fns=reward_fns,
         metric_fns=metric_fns,
         data_shuffle_seed=data_shuffle_seed,
+        diffusion_logits_fn=diffusion_logits_fn,
     )
diff --git a/tunix/rl/grpo/grpo_learner.py b/tunix/rl/grpo/grpo_learner.py
index 7c2ebaaee..278676994 100644
--- a/tunix/rl/grpo/grpo_learner.py
+++ b/tunix/rl/grpo/grpo_learner.py
@@ -24,9 +24,11 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from tunix.rl import algo_core  # pylint: disable=unused-import
+from tunix.diffusion import interfaces as diffusion_interfaces
+from tunix.diffusion import types as diffusion_types
 from tunix.generate import utils
 from tunix.perf.experimental import constants as perf_constants
+from tunix.rl import algo_core  # pylint: disable=unused-import
 from tunix.rl import algorithm_config as algo_config_lib
 from tunix.rl import common
 from tunix.rl import function_registry
@@ -40,7 +42,102 @@
 
 @flax.struct.dataclass(frozen=True)
 class TrainExample(common.TrainExample):
-  pass
+  diffusion_batch: diffusion_types.DiffusionTokenBatch | None = None
+
+
+def _assert_prepared_array_matches(
+    name: str,
+    actual: jax.Array | np.ndarray,
+    expected: np.ndarray,
+    comparison_mask: np.ndarray | None = None,
+) -> None:
+  """Checks a possibly distributed prepared array against host rollout data."""
+
+  if tuple(actual.shape) != tuple(expected.shape):
+    raise ValueError(
+        f"diffusion_batch {name} shape {tuple(actual.shape)} does not match "
+        f"GRPO shape {tuple(expected.shape)}"
+    )
+  if isinstance(actual, jax.Array) and not actual.is_fully_addressable:
+    local_values = [
+        (shard.index, np.asarray(shard.data))
+        for shard in actual.addressable_shards
+    ]
+  else:
+    local_values = [(Ellipsis, np.asarray(actual))]
+  for index, local_actual in local_values:
+    local_expected = expected[index]
+    local_mask = None if comparison_mask is None else comparison_mask[index]
+    if local_mask is not None:
+      local_actual = local_actual[local_mask]
+      local_expected = local_expected[local_mask]
+    if not np.array_equal(local_actual, local_expected):
+      raise ValueError(
+          f"diffusion_batch {name} must exactly match the GRPO rollout"
+      )
+
+
+def _validate_diffusion_rollout_batch(
+    batch: diffusion_types.DiffusionTokenBatch,
+    visible_completion_ids: list[np.ndarray],
+) -> diffusion_types.DiffusionTokenBatch:
+  """Validates that visible completions are prefixes of the full trace."""
+
+  batch = batch.validate()
+  batch_size, trace_length = batch.target_ids.shape
+  if batch_size != len(visible_completion_ids):
+    raise ValueError(
+        "diffusion_batch batch size must match the number of rollout "
+        "completions"
+    )
+  visible_ids = np.zeros((batch_size, trace_length), dtype=np.int64)
+  visible_mask = np.zeros((batch_size, trace_length), dtype=bool)
+  for index, completion_ids in enumerate(visible_completion_ids):
+    completion_ids = np.asarray(completion_ids)
+    if completion_ids.ndim != 1 or completion_ids.shape[0] > trace_length:
+      raise ValueError(
+          "each visible diffusion completion must be one-dimensional and no "
+          "longer than the prepared trace"
+      )
+    visible_ids[index, : completion_ids.shape[0]] = completion_ids
+    visible_mask[index, : completion_ids.shape[0]] = True
+  _assert_prepared_array_matches(
+      "target_ids",
+      batch.target_ids,
+      visible_ids,
+      comparison_mask=visible_mask,
+  )
+  _assert_prepared_array_matches(
+      "active loss_weights",
+      batch.loss_weights != 0,
+      np.ones_like(visible_mask),
+      comparison_mask=visible_mask,
+  )
+  return batch
+
+
+def _stack_rollout_logps(
+    rollout_logps: list[np.ndarray], batch_size: int, trace_length: int
+) -> jax.Array:
+  """Stacks exact full-trace sampler action logps for parity diagnostics."""
+
+  if len(rollout_logps) != batch_size:
+    raise ValueError(
+        "diffusion rollout logprobs must contain one array per completion"
+    )
+  stacked_logps = []
+  for index, logps in enumerate(rollout_logps):
+    logps = np.asarray(logps, dtype=np.float32)
+    if logps.ndim != 1 or logps.shape[0] != trace_length:
+      raise ValueError(
+          "diffusion rollout logprobs must cover the full prepared trace; "
+          f"row {index} has shape {logps.shape} and trace length "
+          f"{trace_length}"
+      )
+    if not np.all(np.isfinite(logps)):
+      raise ValueError("diffusion rollout logprobs must be finite")
+    stacked_logps.append(logps)
+  return jnp.asarray(np.stack(stacked_logps))
 
 
 @dataclasses.dataclass(kw_only=True)
@@ -132,6 +229,7 @@ def __init__(
       reward_fns: RewardFn | List[RewardFn],
       metric_fns: Sequence[MetricFn] | None = None,
       data_shuffle_seed: int | None = None,
+      diffusion_logits_fn: diffusion_interfaces.DiffusionLogitsFn | None = None,
   ):
     """Initializes the `GRPOTrainer`.
 
@@ -157,6 +255,10 @@ def __init__(
            ...       "prompt_min_len": (min(len(p) for p in prompts), np.min),
            ...       # ... }
       data_shuffle_seed: The seed used to shuffle the training data.
+      diffusion_logits_fn: Optional target-aligned scorer for prepared diffusion
+        rollouts. When provided, every rollout must return a matching
+        ``diffusion_batch`` and all live, reference, and old-policy scores use
+        this scorer.
     """  # fmt: skip
     super().__init__(
         rl_cluster=rl_cluster,
@@ -169,6 +271,18 @@ def __init__(
     self.algo_config.temperature = self.rl_cluster.get_rollout_config(  # pyrefly: ignore[missing-attribute]
         mode=rl_cluster_lib.Mode.TRAIN
     ).temperature
+    self.diffusion_logits_fn = diffusion_logits_fn
+    if self.diffusion_logits_fn is not None:
+      if self.algo_config.policy_loss_fn != "grpo":
+        raise ValueError(
+            "diffusion policy scoring is supported only by the GRPO policy "
+            "loss, not PPO"
+        )
+      if (
+          self.rl_cluster.cluster_config.training_config.max_seq_token_per_tpu
+          is not None
+      ):
+        raise ValueError("diffusion GRPO does not support sequence packing")
 
     policy_loss_fn = function_registry.get_policy_loss_fn(
         self.algo_config.policy_loss_fn
@@ -183,6 +297,7 @@ def __init__(
         pad_id=self.rl_cluster.rollout.pad_id(),
         eos_id=self.rl_cluster.rollout.eos_id(),
         compute_logps_chunk_size=self.rl_cluster.cluster_config.training_config.compute_logps_chunk_size,
+        diffusion_logits_fn=self.diffusion_logits_fn,
     )
 
     self.rl_cluster.actor_trainer.with_loss_fn(
@@ -224,7 +339,9 @@ def _generate_and_compute_advantage(
     if isinstance(rollout_config, dict):
       rollout_config = rollout_config[mode]
 
-    training_input["prompts"] = list(training_input["prompts"])  # pyrefly: ignore[bad-argument-type]
+    training_input["prompts"] = list(
+        training_input["prompts"]
+    )  # pyrefly: ignore[bad-argument-type]
     pad_value = self.rl_cluster.rollout.pad_id()
     eos_value = self.rl_cluster.rollout.eos_id()
 
@@ -233,34 +350,64 @@ def _generate_and_compute_advantage(
         perf_constants.STEP: self.rl_cluster.global_steps,
     }
 
+    if (
+        self.diffusion_logits_fn is not None
+        and self.algo_config.num_iterations > 1
+        and not self.should_sync_weights
+    ):
+      self.rl_cluster.snapshot_anchor_policy()
+
     rollout_output = self.rl_cluster.generate(
         prompts=training_input["prompts"],
         mode=mode,
         micro_batch_size=(
-            self._rollout_micro_batch_size * self.algo_config.num_generations  # pyrefly: ignore[unsupported-operation]
+            self._rollout_micro_batch_size
+            * self.algo_config.num_generations  # pyrefly: ignore[unsupported-operation]
         ),
         trace_tags=perf_tags,
     )
-    padded_completion_ids = np.array([
-        utils.pad_to_length(
-            completion_ids,
-            target_length=rollout_config.max_tokens_to_generate,
-            pad_value=pad_value,
-            left=False,
-        )
-        for completion_ids in rollout_output.tokens
-    ])
     prompt_ids = jnp.array(rollout_output.left_padded_prompt_tokens)
-
-    # Assemble masks
     prompt_mask = prompt_ids != pad_value
-    completion_mask = np.not_equal(padded_completion_ids, pad_value)
-
-    # Convert completion_ids and completion_mask to jax arrays
-    jax_completion_ids = jnp.array(padded_completion_ids)
-    jax_completion_mask = jnp.array(completion_mask)
+    diffusion_batch = rollout_output.diffusion_batch
+    if self.diffusion_logits_fn is None:
+      if diffusion_batch is not None:
+        raise ValueError(
+            "rollout returned diffusion_batch without a diffusion_logits_fn; "
+            "refusing autoregressive fallback"
+        )
+      padded_completion_ids = np.array([
+          utils.pad_to_length(
+              completion_ids,
+              target_length=rollout_config.max_tokens_to_generate,
+              pad_value=pad_value,
+              left=False,
+          )
+          for completion_ids in rollout_output.tokens
+      ])
+      completion_mask = np.not_equal(padded_completion_ids, pad_value)
+      jax_completion_ids = jnp.array(padded_completion_ids)
+      jax_completion_mask = jnp.array(completion_mask)
+      visible_completion_lengths = completion_mask.sum(axis=-1)
+    else:
+      if diffusion_batch is None:
+        raise ValueError(
+            "diffusion_logits_fn requires every rollout to return "
+            "diffusion_batch"
+        )
+      diffusion_batch = _validate_diffusion_rollout_batch(
+          diffusion_batch,
+          rollout_output.tokens,
+      )
+      jax_completion_ids = jnp.asarray(diffusion_batch.target_ids)
+      jax_completion_mask = jnp.asarray(diffusion_batch.loss_weights)
+      visible_completion_lengths = np.asarray(
+          [len(tokens) for tokens in rollout_output.tokens], dtype=np.int32
+      )
 
-    if self.algo_config.beta != 0.0:
+    needs_reference = self.algo_config.beta != 0.0
+    if self.diffusion_logits_fn is not None and self.algo_config.beta is None:
+      needs_reference = False
+    if needs_reference:
       devices = self.rl_cluster.r2m[rl_cluster_lib.Role.REFERENCE].devices
       # TODO(yangmu): use function decorator to trace this part, same below.
       with self.rl_cluster.perf.span(
@@ -268,16 +415,27 @@ def _generate_and_compute_advantage(
       ) as interval, self.rl_cluster.perf_v2.span(
           perf_constants.REFERENCE_INFERENCE, devices, tags=perf_tags
       ) as interval_v2:
-        ref_per_token_logps = self.rl_cluster.get_ref_per_token_logps(
-            prompt_tokens=prompt_ids,
-            completion_tokens=jax_completion_ids,
-            pad_id=pad_value,
-            eos_id=eos_value,
-            micro_batch_size=(
-                self._compute_logps_micro_batch_size  # pyrefly: ignore[unsupported-operation]
-                * self.algo_config.num_generations
-            ),
+        compute_logps_micro_batch_size = (
+            self._compute_logps_micro_batch_size  # pyrefly: ignore[unsupported-operation]
+            * self.algo_config.num_generations
         )
+        if self.diffusion_logits_fn is None:
+          ref_per_token_logps = self.rl_cluster.get_ref_per_token_logps(
+              prompt_tokens=prompt_ids,
+              completion_tokens=jax_completion_ids,
+              pad_id=pad_value,
+              eos_id=eos_value,
+              micro_batch_size=compute_logps_micro_batch_size,
+          )
+        else:
+          ref_per_token_logps = (
+              self.rl_cluster.get_ref_diffusion_per_token_logps(
+                  batch=diffusion_batch,
+                  logits_fn=self.diffusion_logits_fn,
+                  micro_batch_size=compute_logps_micro_batch_size,
+                  temperature=self.algo_config.temperature,
+              )
+          )
         interval.device_end([ref_per_token_logps])
         interval_v2.async_end([ref_per_token_logps])
     else:
@@ -289,19 +447,43 @@ def _generate_and_compute_advantage(
       ) as interval, self.rl_cluster.perf_v2.span(
           perf_constants.OLD_ACTOR_INFERENCE, devices, tags=perf_tags
       ) as interval_v2:
-        old_per_token_logps = self.rl_cluster.get_old_per_token_logps(
-            prompt_tokens=prompt_ids,
-            completion_tokens=jax_completion_ids,
-            micro_batch_size=(
-                self._compute_logps_micro_batch_size  # pyrefly: ignore[unsupported-operation]
-                * self.algo_config.num_generations
-            ),
+        compute_logps_micro_batch_size = (
+            self._compute_logps_micro_batch_size  # pyrefly: ignore[unsupported-operation]
+            * self.algo_config.num_generations
         )
+        if self.diffusion_logits_fn is None:
+          old_per_token_logps = self.rl_cluster.get_old_per_token_logps(
+              prompt_tokens=prompt_ids,
+              completion_tokens=jax_completion_ids,
+              micro_batch_size=compute_logps_micro_batch_size,
+          )
+        else:
+          old_per_token_logps = (
+              self.rl_cluster.get_anchor_diffusion_per_token_logps(
+                  batch=diffusion_batch,
+                  logits_fn=self.diffusion_logits_fn,
+                  micro_batch_size=compute_logps_micro_batch_size,
+                  temperature=self.algo_config.temperature,
+              )
+          )
         interval.device_end([old_per_token_logps])
         interval_v2.async_end([old_per_token_logps])
     else:
       old_per_token_logps = None
 
+    if (
+        self.diffusion_logits_fn is not None
+        and old_per_token_logps is not None
+        and rollout_output.logprobs is not None
+    ):
+      rollout_logps = _stack_rollout_logps(
+          rollout_output.logprobs,
+          jax_completion_ids.shape[0],
+          jax_completion_ids.shape[1],
+      )
+    else:
+      rollout_logps = None
+
     with self.rl_cluster.perf.span(
         "advantage_computation"
     ), self.rl_cluster.perf_v2.span(
@@ -312,7 +494,9 @@ def _generate_and_compute_advantage(
           prompts=training_input["prompts"],
           completions=rollout_output.text,
           mode=mode,
-          **{k: v for k, v in training_input.items() if k != "prompts"},  # pyrefly: ignore[bad-argument-type]
+          **{
+              k: v for k, v in training_input.items() if k != "prompts"
+          },  # pyrefly: ignore[bad-argument-type]
       )
       advantage_estimator = function_registry.get_advantage_estimator(
           self.algo_config.advantage_estimator
@@ -332,7 +516,7 @@ def _generate_and_compute_advantage(
     )
 
     # Log completion lengths.
-    agg_completion_mask = completion_mask.sum(axis=-1)
+    agg_completion_mask = visible_completion_lengths
     self.rl_cluster.buffer_metrics(
         {
             "completions/mean_length": (
@@ -362,7 +546,7 @@ def _generate_and_compute_advantage(
       )
       self.rl_cluster.buffer_metrics(user_defined_metric, mode=mode)
 
-    return TrainExample(
+    train_example = TrainExample(
         prompt_ids=prompt_ids,
         prompt_mask=prompt_mask,
         completion_ids=jax_completion_ids,
@@ -370,7 +554,38 @@ def _generate_and_compute_advantage(
         ref_per_token_logps=ref_per_token_logps,
         advantages=jax.device_put(advantages),
         old_per_token_logps=old_per_token_logps,
+        diffusion_batch=diffusion_batch,
     )
+    if self.diffusion_logits_fn is not None:
+      train_example, rollout_logps = self.rl_cluster.place_pytree_on_role(
+          (train_example, rollout_logps), rl_cluster_lib.Role.ACTOR
+      )
+      if rollout_logps is not None:
+        parity_error = jnp.abs(
+            rollout_logps - train_example.old_per_token_logps
+        )
+        active_count = jnp.maximum(jnp.sum(train_example.completion_mask), 1)
+        self.rl_cluster.buffer_metrics(
+            {
+                "diffusion/rollout_anchor_logp_abs_mean": (
+                    jnp.sum(parity_error * train_example.completion_mask)
+                    / active_count,
+                    np.mean,
+                ),
+                "diffusion/rollout_anchor_logp_abs_max": (
+                    jnp.max(
+                        jnp.where(
+                            train_example.completion_mask,
+                            parity_error,
+                            0.0,
+                        )
+                    ),
+                    np.max,
+                ),
+            },
+            mode=mode,
+        )
+    return train_example
 
   def _compute_trajectory_ids(
       self, example: TrainingInputT, steps: int
@@ -388,7 +603,9 @@ def _compute_trajectory_ids(
     Returns:
       A list of trajectory IDs, one for each prompt in the batch.
     """
-    batch_size = len(example["prompts"]) // self.algo_config.num_generations  # pyrefly: ignore[bad-argument-type]
+    batch_size = (
+        len(example["prompts"]) // self.algo_config.num_generations
+    )  # pyrefly: ignore[bad-argument-type]
     row_offset = steps * batch_size
     row_offsets = np.repeat(
         np.arange(row_offset, row_offset + batch_size),
diff --git a/tunix/rl/ppo/ppo_learner.py b/tunix/rl/ppo/ppo_learner.py
index d46e5e167..6a065d30e 100644
--- a/tunix/rl/ppo/ppo_learner.py
+++ b/tunix/rl/ppo/ppo_learner.py
@@ -32,7 +32,6 @@
 from tunix.rl import rl_cluster as rl_cluster_lib
 from tunix.rl import rl_learner
 
-
 TrainingInputT = rl_learner.TrainingInputT
 RewardFn = rl_learner.RewardFn
 MetricFn = rl_learner.MetricFn
@@ -269,6 +268,11 @@ def _generate_and_compute_advantage(
         prompts=training_input["prompts"],  # pyrefly: ignore[bad-argument-type]
         micro_batch_size=self._rollout_micro_batch_size,
     )
+    if rollout_output.diffusion_batch is not None:
+      raise ValueError(
+          "prepared diffusion rollouts are not supported by PPO; use the "
+          "standard GRPO family with a diffusion_logits_fn"
+      )
     padded_completion_ids = np.array([
         utils.pad_to_length(
             completion_ids,
@@ -343,10 +347,14 @@ def _generate_and_compute_advantage(
       last_token_scores = jax.device_get(jax_last_token_scores)
     else:
       last_token_scores = self._compute_rewards(
-          prompts=training_input["prompts"],  # pyrefly: ignore[bad-argument-type]
+          prompts=training_input[
+              "prompts"
+          ],  # pyrefly: ignore[bad-argument-type]
           completions=rollout_output.text,
           mode=mode,
-          **{k: v for k, v in training_input.items() if k != "prompts"},  # pyrefly: ignore[bad-argument-type]
+          **{
+              k: v for k, v in training_input.items() if k != "prompts"
+          },  # pyrefly: ignore[bad-argument-type]
       )
       jax_last_token_scores = jax.device_put(last_token_scores)
 
@@ -516,7 +524,9 @@ def _compute_trajectory_ids(
     Returns:
       A list of trajectory IDs, one for each prompt in the batch.
     """
-    batch_size = len(example["prompts"]) // self._num_generations()  # pyrefly: ignore[bad-argument-type]
+    batch_size = (
+        len(example["prompts"]) // self._num_generations()
+    )  # pyrefly: ignore[bad-argument-type]
     row_offset = steps * batch_size
     row_offsets = np.arange(row_offset, row_offset + batch_size)
     return row_offsets.astype(str).tolist()
diff --git a/tunix/rl/rl_cluster.py b/tunix/rl/rl_cluster.py
index 561a57e2e..e1860a604 100644
--- a/tunix/rl/rl_cluster.py
+++ b/tunix/rl/rl_cluster.py
@@ -22,6 +22,7 @@
 import functools
 import gc
 import itertools
+import math
 import operator
 import os
 from typing import Any, Callable, Mapping
@@ -38,6 +39,7 @@
 import jaxtyping
 import numpy as np
 import optax
+from tunix.diffusion import interfaces as diffusion_interfaces
 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.
@@ -46,6 +48,7 @@
 from tunix.perf.experimental import constants as perf_constants
 from tunix.perf.experimental import tracer as perf_tracer_v2
 from tunix.rl import common
+from tunix.rl import diffusion as diffusion_lib
 from tunix.rl import reshard
 from tunix.rl import trainer as rl_trainer
 from tunix.rl import utils as rl_utils
@@ -100,6 +103,8 @@ def _merge_diffusion_batches(
     )
 
   def concatenate(*arrays):
+    if all(isinstance(array, np.ndarray) for array in arrays):
+      return np.concatenate(arrays, axis=0)
     return jnp.concatenate([jnp.asarray(array) for array in arrays], axis=0)
 
   return diffusion_types.DiffusionTokenBatch.create(
@@ -116,6 +121,84 @@ def concatenate(*arrays):
   )
 
 
+def _is_custom_rollout_engine(rollout_engine: Any) -> bool:
+  """Returns whether a rollout class consumes the supplied model directly."""
+
+  return (
+      isinstance(rollout_engine, type)
+      and issubclass(rollout_engine, base_rollout.BaseRollout)
+  ) or (
+      isinstance(rollout_engine, functools.partial)
+      and isinstance(rollout_engine.func, type)
+      and issubclass(rollout_engine.func, base_rollout.BaseRollout)
+  )
+
+
+def _slice_diffusion_batch(
+    batch: diffusion_types.DiffusionTokenBatch, batch_slice: slice
+) -> diffusion_types.DiffusionTokenBatch:
+  """Slices every prepared diffusion input along its batch dimension."""
+
+  return jax.tree.map(lambda value: value[batch_slice], batch)
+
+
+@nnx.jit(static_argnames=("logits_fn", "temperature"))
+def _compute_detached_diffusion_logps(
+    model: nnx.Module,
+    batch: diffusion_types.DiffusionTokenBatch,
+    logits_fn: diffusion_interfaces.DiffusionLogitsFn,
+    temperature: float,
+) -> jax.Array:
+  """Compiles detached policy scoring for reference and anchor roles."""
+
+  return diffusion_lib.compute_diffusion_per_token_logps(
+      model,
+      batch,
+      logits_fn,
+      temperature=temperature,
+      stop_gradient=True,
+  )
+
+
+def _place_pytree_on_mesh(
+    pytree: Any,
+    mesh: Mesh,
+    data_sharding_axis: tuple[str, ...],
+) -> Any:
+  """Places every array leaf on one role mesh, including cross-mesh arrays."""
+
+  if mesh.empty:
+    return pytree
+  has_non_addressable_array = any(
+      isinstance(leaf, jax.Array) and not leaf.is_fully_addressable
+      for leaf in jax.tree.leaves(pytree)
+  )
+  if not has_non_addressable_array:
+    return sharding_utils.shard_input(pytree, data_sharding_axis)
+
+  pspec = jax.sharding.PartitionSpec(*data_sharding_axis)
+  dst_shardings = jax.tree.map(
+      lambda value: sharding_utils.get_sharding(value, mesh, pspec), pytree
+  )
+
+  def place_leaf(value, dst_sharding):
+    if isinstance(value, jax.Array) and not value.is_fully_addressable:
+      return reshard.reshard_pytree(value, dst_sharding)
+    return jax.device_put(value, dst_sharding)
+
+  return jax.tree.map(place_leaf, pytree, dst_shardings)
+
+
+def _shard_diffusion_batch(
+    batch: diffusion_types.DiffusionTokenBatch,
+    mesh: Mesh,
+    data_sharding_axis: tuple[str, ...],
+) -> diffusion_types.DiffusionTokenBatch:
+  """Places every prepared leaf on the requested role mesh."""
+
+  return _place_pytree_on_mesh(batch, mesh, data_sharding_axis)
+
+
 class Mode(enum.Enum):
   """Mode of RolloutConfig."""
 
@@ -169,9 +252,44 @@ class RLTrainingConfig(peft_trainer.TrainingConfig):
   rollout_micro_batch_size: int | None = None
   compute_logps_micro_batch_size: int | None = None
   compute_logps_chunk_size: int = 0
+  checkpoint_metadata: Mapping[str, Any] = dataclasses.field(
+      default_factory=dict
+  )
 
   def __post_init__(self):
     """Validates the configuration after initialization."""
+    if not isinstance(self.checkpoint_metadata, Mapping):
+      raise TypeError("checkpoint_metadata must be a mapping")
+    self.checkpoint_metadata = dict(self.checkpoint_metadata)
+    reserved_keys = {"global_step", "role"}.intersection(
+        self.checkpoint_metadata
+    )
+    if reserved_keys:
+      raise ValueError(
+          "checkpoint_metadata cannot override reserved keys: "
+          f"{sorted(reserved_keys)}"
+      )
+
+    def validate_metadata_value(name: str, value: Any) -> None:
+      if isinstance(value, float) and not math.isfinite(value):
+        raise ValueError(
+            f"checkpoint_metadata[{name!r}] must be a finite JSON value"
+        )
+      if value is None or isinstance(value, (str, bool, int, float)):
+        return
+      if isinstance(value, tuple):
+        for index, item in enumerate(value):
+          validate_metadata_value(f"{name}[{index}]", item)
+        return
+      raise TypeError(
+          f"checkpoint_metadata[{name!r}] must be a JSON scalar or tuple"
+      )
+
+    for key, value in self.checkpoint_metadata.items():
+      if not isinstance(key, str):
+        raise TypeError("checkpoint_metadata keys must be strings")
+      validate_metadata_value(key, value)
+
     for name in [
         "mini_batch_size",
         "train_micro_batch_size",
@@ -263,7 +381,7 @@ def __init__(
     self._anchor_policy_state = None
 
     self._default_memory_kind = jax.devices()[0].default_memory().kind
-    self.train_actor = self._load_model(actor, self.r2m[Role.ACTOR])
+    self.train_actor = self._load_model(actor, Role.ACTOR)
 
     if self.cluster_config.rollout_config is None:
       raise ValueError("`cluster_config.rollout_config` cannot be None.")
@@ -277,7 +395,9 @@ def __init__(
 
     if Role.ROLLOUT in self._backbone_sharing_map[Role.ACTOR]:
       self.rollout_actor = self.train_actor
-    elif self.cluster_config.rollout_engine == "vanilla":
+    elif self.cluster_config.rollout_engine == "vanilla" or (
+        _is_custom_rollout_engine(self.cluster_config.rollout_engine)
+    ):
       rollout_data_type = (
           self.cluster_config.rollout_config[Mode.TRAIN].data_type
           if isinstance(self.cluster_config.rollout_config, dict)
@@ -285,7 +405,7 @@ def __init__(
       )
       self.rollout_actor = self._load_model(
           actor,
-          self.r2m[Role.ROLLOUT],
+          Role.ROLLOUT,
           rollout_data_type,
       )
     else:
@@ -293,7 +413,7 @@ def __init__(
       self.rollout_actor = self.train_actor
 
     if reference:
-      self.reference = self._load_model(reference, self.r2m[Role.REFERENCE])
+      self.reference = self._load_model(reference, Role.REFERENCE)
       if Role.REFERENCE in self._backbone_sharing_map[Role.ACTOR]:
         if not rl_utils.is_sharing_backbone(self.reference, self.train_actor):
           logging.warning(
@@ -303,15 +423,11 @@ def __init__(
           )
     else:
       self.reference = None
-    self.critic = (
-        self._load_model(critic, self.r2m[Role.CRITIC]) if critic else None
-    )
+    self.critic = self._load_model(critic, Role.CRITIC) if critic else None
     if Role.CRITIC in self._backbone_sharing_map[Role.ACTOR]:
       critic_state = nnx.state(self.train_actor, filterlib.Not(nnx.LoRAParam))
       nnx.update(self.critic, critic_state)
-    self.reward = (
-        self._load_model(reward, self.r2m[Role.REWARD]) if reward else None
-    )
+    self.reward = self._load_model(reward, Role.REWARD) if reward else None
 
     self.tokenizer = tokenizer_adapter.TokenizerAdapter(tokenizer)
     self._rl_metrics_logger = metrics_logger.MetricsLogger(
@@ -364,7 +480,7 @@ def _init_backbone_sharing_map(
   def _load_model(
       self,
       model_or_path: ModelOrPath,
-      mesh: Mesh,
+      role: Role,
       data_type: jnp.dtype | None = None,
   ) -> nnx.Module:
     """Loads model with given mesh to the given memory_kind.
@@ -374,12 +490,13 @@ def _load_model(
 
     Args:
       model_or_path: either a nnx.Module or a path to a model.
-      mesh: the mesh to load the model on.
+      role: the role whose mesh and logical axis rules should be used.
       data_type: optional data type to cast the model parameters to.
 
     Returns:
       The model loaded on the given mesh.
     """
+    mesh = self.r2m[role]
     if isinstance(model_or_path, nnx.Module):
       model_mesh = rl_utils.get_pytree_mesh_info(nnx.state(model_or_path))
       original_shardings = jax.tree_util.tree_map(
@@ -395,6 +512,17 @@ def _load_model(
       if not mesh.empty and model_mesh != mesh:
         logging.warning("Resharding model from %s to %s", model_mesh, mesh)
         graph, state = nnx.split(model_or_path)
+        logical_axis_rules = (
+            self.cluster_config.role_to_logical_axis_rule or {}
+        ).get(role)
+        physical_specs = nnx.get_partition_spec(state)
+        if logical_axis_rules is not None:
+          physical_specs = jax.tree_util.tree_map(
+              lambda spec: nn_partitioning.logical_to_mesh_axes(
+                  spec, logical_axis_rules
+              ),
+              physical_specs,
+          )
         dst_shardings = jax.tree_util.tree_map(
             lambda x: jax.sharding.NamedSharding(
                 mesh,
@@ -403,7 +531,7 @@ def _load_model(
                 if is_on_device
                 else "pinned_host",
             ),
-            nnx.get_partition_spec(state),
+            physical_specs,
         )
         if data_type and data_type != jax.tree.leaves(state)[0].dtype:
           tmp_state = jax.tree.map(lambda x: x.astype(data_type), state)
@@ -521,18 +649,7 @@ def _init_cluster(self):
           mesh=self.r2m[Role.ROLLOUT],
           rollout_config=loaded_sglang_jax_config,
       )
-    elif (
-        isinstance(self.cluster_config.rollout_engine, type)
-        and issubclass(
-            self.cluster_config.rollout_engine, base_rollout.BaseRollout
-        )
-    ) or (
-        isinstance(self.cluster_config.rollout_engine, functools.partial)
-        and issubclass(
-            self.cluster_config.rollout_engine.func,  # pyrefly: ignore[bad-argument-type]
-            base_rollout.BaseRollout,
-        )
-    ):
+    elif _is_custom_rollout_engine(self.cluster_config.rollout_engine):
       if isinstance(self.cluster_config.rollout_config, dict):
         loaded_config = self.cluster_config.rollout_config[Mode.TRAIN]
       else:
@@ -614,6 +731,7 @@ def _init_cluster(self):
             optimizer=self.cluster_config.training_config.critic_optimizer,  # pyrefly: ignore[bad-argument-type]
             training_config=critic_config,
             custom_checkpoint_metadata_fn=lambda: {
+                **self.cluster_config.training_config.checkpoint_metadata,
                 "global_step": self.global_steps + 1,
                 "role": Role.CRITIC.value,
             },  # offset by 1 since global_step is incremented after the training loop in rl_learner. # pylint: disable=line-too-long
@@ -638,6 +756,7 @@ def _init_cluster(self):
           optimizer=self.cluster_config.training_config.actor_optimizer,
           training_config=actor_config,
           custom_checkpoint_metadata_fn=lambda: {
+              **self.cluster_config.training_config.checkpoint_metadata,
               "global_step": self.global_steps + 1,
               "role": Role.ACTOR.value,
           },  # offset by 1 since global_step is incremented after the training loop in rl_learner. # pylint: disable=line-too-long
@@ -648,9 +767,7 @@ def _init_cluster(self):
     del self.rollout_actor
     del self.train_actor
     self._maybe_offload_model_to_cpu(self.actor_trainer.model, Role.ACTOR)
-    self._anchor_policy_state = rl_utils.put_params_on_memory_kind(
-        nnx.state(self.actor_trainer.model), "pinned_host"
-    )
+    self.snapshot_anchor_policy()
 
   def _propagate_backbone_sharing_map(self):
     """Propagates backbone sharing map."""
@@ -987,6 +1104,12 @@ def generate(
             max_tokens_to_generate=max_generation_steps,
         )
 
+      set_generation_context = getattr(
+          self.rollout, "set_generation_context", None
+      )
+      if set_generation_context is not None:
+        set_generation_context(global_step=self.global_steps, mode=mode)
+
       perf_tags = {
           perf_constants.ROLE: Role.ROLLOUT.value,
       }
@@ -1016,7 +1139,9 @@ def generate(
     logprobs = None
     if outputs[0].logprobs is not None:
       logprobs = list(
-          itertools.chain.from_iterable(out.logprobs for out in outputs)  # pyrefly: ignore[bad-argument-type]
+          itertools.chain.from_iterable(
+              out.logprobs for out in outputs
+          )  # pyrefly: ignore[bad-argument-type]
       )
 
     logits = None
@@ -1088,6 +1213,50 @@ def get_ref_per_token_logps(
       )
       return ref_per_token_logps
 
+  def get_ref_diffusion_per_token_logps(
+      self,
+      batch: diffusion_types.DiffusionTokenBatch,
+      logits_fn: diffusion_interfaces.DiffusionLogitsFn,
+      micro_batch_size: int | None = None,
+      temperature: float | None = None,
+  ) -> jax.Array:
+    """Scores a prepared batch with the detached reference policy."""
+
+    batch = batch.validate()
+    batch_size = batch.target_ids.shape[0]
+    if batch_size == 0:
+      raise ValueError(
+          "Cannot get reference diffusion log probabilities from an empty"
+          " batch."
+      )
+    micro_batch_size = micro_batch_size or batch_size
+    if temperature is None:
+      temperature = self.get_rollout_config(mode=Mode.TRAIN).temperature
+
+    with self._get_mesh_and_logical_axis_rules_cm(Role.REFERENCE):
+      dest_batch = _shard_diffusion_batch(
+          batch,
+          self.r2m[Role.REFERENCE],
+          self.cluster_config.training_config.data_sharding_axis,
+      )
+      model = self.inference_worker.get_model("reference")
+      self._maybe_load_model_from_cpu(model, Role.REFERENCE)
+      outs = []
+      for batch_slice in rl_utils.chunk_slices_by_size(
+          stop=batch_size, step=micro_batch_size
+      ):
+        outs.append(
+            _compute_detached_diffusion_logps(
+                model,
+                _slice_diffusion_batch(dest_batch, batch_slice),
+                logits_fn,
+                temperature=temperature,
+            )
+        )
+      ref_per_token_logps = jnp.concatenate(outs, axis=0)
+      self._maybe_offload_model_to_cpu(model, Role.REFERENCE)
+      return ref_per_token_logps
+
   def get_old_per_token_logps(
       self,
       prompt_tokens: jax.Array,
@@ -1209,7 +1378,93 @@ def get_actor_per_token_logps(
         )
       return actor_per_token_logps
 
-  def sync_weights(self):
+  def get_anchor_diffusion_per_token_logps(
+      self,
+      batch: diffusion_types.DiffusionTokenBatch,
+      logits_fn: diffusion_interfaces.DiffusionLogitsFn,
+      micro_batch_size: int | None = None,
+      temperature: float | None = None,
+  ) -> jax.Array:
+    """Scores a prepared batch with the start-of-step actor snapshot."""
+
+    batch = batch.validate()
+    batch_size = batch.target_ids.shape[0]
+    if batch_size == 0:
+      raise ValueError(
+          "Cannot get anchor diffusion log probabilities from an empty batch."
+      )
+    if self._anchor_policy_state is None:
+      raise ValueError(
+          "Anchor policy state is not initialized. Please run `sync_weights`"
+          " first."
+      )
+    micro_batch_size = micro_batch_size or batch_size
+    if temperature is None:
+      temperature = self.get_rollout_config(mode=Mode.TRAIN).temperature
+
+    with self._get_mesh_and_logical_axis_rules_cm(Role.ACTOR):
+      dest_batch = _shard_diffusion_batch(
+          batch,
+          self.r2m[Role.ACTOR],
+          self.cluster_config.training_config.data_sharding_axis,
+      )
+      actor_state_on_device = self._is_state_on_device(
+          nnx.state(self.actor_trainer.model)
+      )
+      if actor_state_on_device and self.cluster_config.offload_to_cpu:
+        self._put_model_on_memory_kind(self.actor_trainer.model, "pinned_host")
+        gc.collect()
+
+      graphdef, _ = nnx.split(self.actor_trainer.model)
+      anchor_on_device = self._is_state_on_device(self._anchor_policy_state)
+      if anchor_on_device:
+        anchor_policy_state = self._anchor_policy_state
+      else:
+        anchor_policy_state = rl_utils.put_params_on_memory_kind(
+            self._anchor_policy_state, self._default_memory_kind
+        )
+      anchor_model = nnx.merge(graphdef, anchor_policy_state)
+      outs = []
+      for batch_slice in rl_utils.chunk_slices_by_size(
+          stop=batch_size, step=micro_batch_size
+      ):
+        outs.append(
+            _compute_detached_diffusion_logps(
+                anchor_model,
+                _slice_diffusion_batch(dest_batch, batch_slice),
+                logits_fn,
+                temperature=temperature,
+            )
+        )
+      anchor_per_token_logps = jnp.concatenate(outs, axis=0)
+      del anchor_model
+      if not anchor_on_device:
+        del anchor_policy_state
+      gc.collect()
+      if actor_state_on_device and self.cluster_config.offload_to_cpu:
+        self._put_model_on_memory_kind(
+            self.actor_trainer.model, self._default_memory_kind
+        )
+      return anchor_per_token_logps
+
+  def snapshot_anchor_policy(self) -> None:
+    """Snapshots the current actor for old-policy scoring within one step."""
+
+    self._anchor_policy_state = rl_utils.put_params_on_memory_kind(
+        nnx.state(self.actor_trainer.model), "pinned_host"
+    )
+
+  def place_pytree_on_role(self, pytree: Any, role: Role) -> Any:
+    """Places a complete input pytree on exactly one role mesh."""
+
+    with self._get_mesh_and_logical_axis_rules_cm(role):
+      return _place_pytree_on_mesh(
+          pytree,
+          self.r2m[role],
+          self.cluster_config.training_config.data_sharding_axis,
+      )
+
+  def sync_weights(self, *, increment_global_steps: bool = True):
     """Syncs the weights of between the sampler model and trainer model."""
     if jax.devices() and jax.default_backend() not in ["tpu", "gpu"]:
       cm = contextlib.ExitStack()
@@ -1227,12 +1482,11 @@ def sync_weights(self):
       self.rollout.update_params(src_filtered_params, filter_types)
       gc.collect()
       # The anchor policy state is snapshotted from actor_trainer.model.
-      self._anchor_policy_state = rl_utils.put_params_on_memory_kind(
-          nnx.state(self.actor_trainer.model), "pinned_host"
-      )
+      self.snapshot_anchor_policy()
 
-    # sync weights marks the end of a full batch, so increment the global steps.
-    self.global_steps += 1
+    if increment_global_steps:
+      # A regular sync marks the end of a full batch.
+      self.global_steps += 1
 
   def get_values(
       self,
diff --git a/tunix/rl/rl_learner.py b/tunix/rl/rl_learner.py
index 805db5068..4291fb356 100644
--- a/tunix/rl/rl_learner.py
+++ b/tunix/rl/rl_learner.py
@@ -36,7 +36,6 @@
 from tunix.rl.queue import data_queue as queue_lib
 from tunix.sft import utils as sft_utils
 
-
 ABC = abc.ABC
 abstractmethod = abc.abstractmethod
 
@@ -118,6 +117,8 @@ def __init__(
             self.rl_cluster.rollout.model(),
         )
     )
+    if self.should_sync_weights:
+      self.rl_cluster.sync_weights(increment_global_steps=False)
 
     # Enable async rollout if trainer and rollout are not on the same mesh.
     # If they do, then doesn't make sense for the interleave because they will
@@ -394,7 +395,9 @@ def _process_and_enqueue_tail():
         ):
           # Fetch one training micro-batch
           example = next(iterator)
-          cur_batch_size = len(example["prompts"])  # pyrefly: ignore[bad-argument-type]
+          cur_batch_size = len(
+              example["prompts"]
+          )  # pyrefly: ignore[bad-argument-type]
 
         # Buffer the fetched micro-batch. We accumulate micro-batches and track
         # their sizes and the total number of samples. This allows us to form a
@@ -527,7 +530,9 @@ def train(
     """Main entry point for the training loop."""
     full_batch_iterator = iter(train_ds)
     first_item = next(full_batch_iterator)
-    full_batch_size = len(first_item["prompts"])  # pyrefly: ignore[bad-argument-type]
+    full_batch_size = len(
+        first_item["prompts"]
+    )  # pyrefly: ignore[bad-argument-type]
     full_batch_iterator = itertools.chain([first_item], full_batch_iterator)
     # Initialize batch sizes.
     mini_batch_size = self._training_config.mini_batch_size or full_batch_size
diff --git a/tunix/rl/trainer.py b/tunix/rl/trainer.py
index 084bcdfb9..560234280 100644
--- a/tunix/rl/trainer.py
+++ b/tunix/rl/trainer.py
@@ -19,11 +19,11 @@
 from flax import nnx
 from jax.typing import ArrayLike  # pylint: disable=g-importing-member
 import optax
-from tunix.sft import peft_trainer
-from typing_extensions import override
 from tunix.perf import trace as perf_trace
 from tunix.perf.experimental import tracer as perf_tracer_lib
+from tunix.sft import peft_trainer
 from tunix.sft.metrics_logger import MetricsLogger  # pylint: disable=unused-import
+from typing_extensions import override
 
 
 class Trainer(peft_trainer.PeftTrainer):
@@ -49,10 +49,38 @@ def __init__(
         perf_tracer,
         perf_tracer_v2,
     )
+    self._validate_restored_checkpoint_metadata(
+        getattr(training_config, "checkpoint_metadata", {})
+    )
     self.rl_metrics_to_log = {}  # Metric name -> key in aux.
     self.tqdm_metrics_to_display = []
     self.custom_checkpoint_metadata_fn = custom_checkpoint_metadata_fn
 
+  def _validate_restored_checkpoint_metadata(
+      self, expected_metadata: dict[str, Any]
+  ) -> None:
+    """Fails closed when a resumed RL objective differs from its checkpoint."""
+
+    if self._train_steps <= 0 or not expected_metadata:
+      return
+
+    def canonicalize(value: Any) -> Any:
+      if isinstance(value, (tuple, list)):
+        return tuple(canonicalize(item) for item in value)
+      return value
+
+    for key, expected_value in expected_metadata.items():
+      if key not in self._restored_custom_metadata:
+        raise ValueError(
+            f"restored checkpoint is missing required metadata key {key!r}"
+        )
+      restored_value = self._restored_custom_metadata[key]
+      if canonicalize(restored_value) != canonicalize(expected_value):
+        raise ValueError(
+            f"restored checkpoint metadata mismatch for {key!r}: expected "
+            f"{expected_value!r}, received {restored_value!r}"
+        )
+
   def with_rl_metrics_to_log(
       self,
       rl_metrics_to_log: dict[str, Callable[[ArrayLike], ArrayLike]],