[tunix] Add prepared diffusion policy scoring - #1748
Open
ethannnnnn wants to merge 6 commits into
Open
Conversation
Define a target-aligned, batch-major diffusion batch contract and typed adapter/scorer protocols without depending on MaxText or a specific training algorithm. Validate shapes and dtypes at construction and scoring boundaries, while preserving JAX pytree, JIT, and sharding compatibility. Tests: 10 diffusion contract tests; pyink/isort; pylint; pyrefly; py_compile.
Accumulate LossOutput gradients as unreduced sums and normalize once by the total denominator across microbatches. Preserve denominator-one behavior for scalar losses and return zero gradients when every weight is zero. Select auxiliary-metric reducers by value type in training and evaluation: globally combine weighted metrics while averaging ordinary scalar metrics. Reject per-key type changes across microbatches and preserve consistent epsilon and minimum-denominator bounds during global reduction. Preserve the dtype selected by each Optax optimizer-state initializer across conditional update and skip branches. This keeps explicit bf16 moments in bf16, retains explicit fp32 moments, and prevents Flax NNX branch-type mismatches without special-casing a particular accumulation count. Tests cover weighted and fractional denominators, zero-weight batches, mixed weighted/plain train and eval metrics, reducer invariants, and a real PeftTrainer + nnx.jit matrix over direct/injected AdamW and gradient accumulation counts 1 and 2. The complete PeftTrainer suite passes 56 tests; the cumulative focused validation passes 119 tests with six optional engine tests deselected. Ruff and git diff checks pass.
Provide a typed PeftTrainer adapter for canonical diffusion batches and target-aligned score functions. Compute weighted float32 cross entropy without autoregressive shifting, sanitize inactive targets, and preserve zero-weight numerical safety. Tests: 17 diffusion contract and SFT tests; 6 focused weighted-gradient tests; pyink, isort, pyrefly, pylint, py_compile, and diff checks.
Define a framework-neutral external-teacher batch contract for freshly prepared student rollouts. Validate the canonical student batch and target-aligned teacher logits without owning model rollout, corruption, or checkpoint behavior. Tests: 3 focused batch-contract tests; included in the 34-test diffusion contract/SFT/OPD suite.
Add forward teacher-to-student KL with temperature scaling, optional target-aligned hard CE, fractional token weights, teacher stop-gradient, inactive-token sanitization, and PeftTrainer wiring for externally prepared fresh rollouts. Tests: 11 focused OPD tests; 34 combined diffusion contract/SFT/OPD tests and 8 weighted trainer regressions passed.
Add a framework-neutral policy scorer for target-aligned DiffusionTokenBatch inputs. The scorer applies the rollout temperature, sanitizes inactive padding, optionally reports finite entropy, and makes gradient ownership explicit. Model-aware integrations remain responsible for preparing the exact action context or denoising trace. Allow rollout engines to attach one optional prepared diffusion batch and merge that pytree across rollout microbatches with strict presence, structure, and batch-size validation. Existing autoregressive RolloutOutput construction and generation remain unchanged when the field is absent. Test plan: - 16 focused diffusion/rollout generation tests passed - 34 existing diffusion CFT/SFT/OPD tests and 6 subtests passed - Pyink 25.12, Pylint 10/10, isort, and git diff --check
ethannnnnn
requested review from
abheesht17,
hgao327,
jiangyangmu,
lc5211,
s-noghabi,
sizhit2,
tianshub and
wang2yn84
as code owners
July 23, 2026 22:43
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
RL rollout output needs an optional, typed diffusion trace that can be merged across microbatches and scored without treating diffusion actions as autoregressive next tokens.
Scope
RolloutOutputto carry one optional prepared diffusion batch.Design
The scorer consumes
DiffusionTokenBatchwhose model inputs already reconstruct the sampled action context. It returns target-aligned action log probabilities and optional entropy while making gradient ownership explicit. Model-aware integrations remain responsible for creating the exact denoising trace.Microbatch merging is atomic: either every rollout shard carries a compatible diffusion batch or the merge fails. No partial metadata is silently discarded.
Compatibility
RolloutOutput.diffusion_batchdefaults to absent. Existing autoregressive rollout construction, merging, and generation are unchanged.Extensibility
The contract can carry traces from MaxText, TPU inference, or another runtime. Learners can consume one scorer interface without knowing how a model reconstructed its diffusion canvas.
Tests
git diff --checkpassed.Known limitations
This PR defines policy scoring and rollout transport only. It does not yet route live/reference/old policy roles through GRPO.
Stack
Depends on the preceding upstream PR: #1747
Tunix block-diffusion design document