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