[tunix] Add exact prepared diffusion policy optimization#1749
Open
ethannnnnn wants to merge 7 commits into
Open
[tunix] Add exact prepared diffusion policy optimization#1749ethannnnnn wants to merge 7 commits into
ethannnnnn wants to merge 7 commits into
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
Teach the GRPO family to optimize prepared block-diffusion rollouts without routing them through autoregressive token scoring. The rollout now carries its full denoising trace as a DiffusionTokenBatch. Live, reference, and start-of-step policies score the same target-aligned states and actions, including weighted objectives. The actor receives the complete prepared example on its own mesh, and multi-iteration updates use a stable old-policy snapshot. Keep the integration opt-in through diffusion_logits_fn. Existing autoregressive GRPO is unchanged, while PPO, agentic RL, sequence packing, and incomplete diffusion metadata fail closed. Checkpoint metadata validates that resumed runs retain the same diffusion objective, and custom rollout models synchronize once at initialization without advancing global_step. Test Plan: - JAX_PLATFORMS=cpu XLA_FLAGS=--xla_force_host_platform_device_count=2 uv run --no-project --with-editable '.[test]' python -m pytest -q tests/rl/diffusion_test.py tests/rl/grpo/grpo_diffusion_test.py --disable-warnings - uv run --no-project --with isort python -m isort --check-only <changed files> - uv run --no-project --with pyink==24.10.1 python -m pyink --check <changed files> - uv run --no-project --with-editable '.[test]' --with pylint python -m pylint --disable=all --enable=E,F <changed files> - python3 -m py_compile <changed Python files> - git diff --check
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
Diffusion GRPO must score sampled actions symmetrically under live, reference, and old policies. Updating only one role or comparing against autoregressive log probabilities produces invalid KL and importance ratios.
Scope
and a detached start-of-step actor snapshot used for old-policy scoring.
global_step.Design
When
diffusion_logits_fnis configured, the GRPO family bypasses autoregressive token scoring and evaluates the full prepared trace for every required policy. The old policy is not a separate model role: it is a snapshot of the actor at the start of the optimization step and is scored on the actor role mesh. Rollout metadata remains host-resident until placement, then the completeTrainExampleis placed on the actor mesh. DAPO forwards the scorer and DrGRPO inherits the same learner path.Checkpoint metadata validates the diffusion objective on resume.
RLClusterexposes an optional generation-context hook so a rollout can derive deterministic train/eval RNG streams from global step without algorithm-specific coupling.Model loading is role-aware. Before constructing each actor, rollout, reference, critic, or reward
NamedSharding,RLClusterconverts the model's logical partition metadata with that role's declared logical-axis rules and physical mesh. Roles without logical-axis rules retain the existing physical-spec path. This is a generic NNX sharding correction and contains no model-specific axis names.Compatibility
The diffusion integration is opt-in. Existing autoregressive GRPO behavior is unchanged, and unsupported algorithms fail during validation instead of receiving diffusion-only arguments. Logical-axis conversion applies generically to RL model loading only when a role declares axis rules.
Extensibility
The learner depends on the prepared scorer, not a MaxText model. Other block-diffusion integrations can provide the same trace contract, and future optimized scorers can replace dense replay without changing the algorithm.
Tests
PeftTrainertests and 63 focused RL/role-sharding tests. Six optionalvLLM/SGL-JAX integration tests were deselected.
fsdp/tpmetadata onto a distinctdata/modelmesh and preserves the existing no-rules resharding path.all three configured actor updates. Losses and rewards were finite,
importance-sampling ratio was 1, clip fraction was 0, and gradient norms
were nonzero. The actor step-3 checkpoint finalized and the downstream model
export was verified. This is a correctness smoke, not a throughput or
model-quality benchmark.
compileall, and
git diff --checkwere clean. Scoped production/test Pylintpassed at 10/10.
Known limitations
PPO, agentic RL, and sequence packing are not supported for prepared diffusion traces. The exact model-aware replay can be compute and memory intensive. The 4x4x4 smoke validates the configured three-update workflow and checkpoint/export handoff, but it does not establish throughput, memory scaling, convergence, or model quality.
Stack
Depends on the preceding upstream PR: #1748
Tunix block-diffusion design document