[tunix] Add weighted prepared diffusion OPD objective#1747
Open
ethannnnnn wants to merge 5 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.
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
Prepared diffusion rollouts need a reusable objective that combines teacher distributions, optional hard targets, explicit fractional token weights, and correct trainer/checkpoint behavior.
Scope
KL(teacher || student)with temperature-squared scaling.PeftTrainerand preserve custom semantic metadata on periodic and final checkpoints.Design
configure_prepared_diffusion_opdconsumes an externally prepared fresh batch. Student logits come from the canonical target-aligned scorer; teacher logits are treated as constants. KL and hard CE share the explicit weights and denominator-awareLossOutputreduction.Trainer metadata hooks default to an empty mapping, while model-aware integrations can record the semantic identity required to reject incompatible resume attempts.
Compatibility
Existing trainers and distillation objectives remain unchanged unless this prepared objective is selected. The metadata hook is empty by default for existing users.
Extensibility
Model integrations own rollout and alignment, so the same Tunix objective can support additional diffusion architectures. A future top-k teacher representation can be introduced behind the prepared-batch boundary.
Tests
PeftTrainerandPeftTrainerV2.Known limitations
Tunix neither performs nor verifies on-policy generation. Dense teacher logits are required, and external artifact references must be pinned by the caller if reproducible resume is required.
Stack
Depends on the preceding upstream PR: #1746
Tunix block-diffusion design document