[tunix] Add prepared diffusion distillation batches#1746
Open
ethannnnnn wants to merge 4 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.
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
On-policy distillation needs a boundary for externally generated student rollouts and immutable teacher targets without rewriting the existing trainer-owned teacher workflow.
Scope
DiffusionDistillationBatchover a canonical student batch and target-aligned dense teacher logits.Design
The contract is deliberately a prepared-data boundary. A model-aware integration generates a fresh rollout from the current student, constructs the student diffusion batch, and scores those exact physical targets with the teacher before entering Tunix.
Compatibility
The existing
DistillationTrainerand dataset-driven teacher path remain unchanged. No rollout, teacher execution, or checkpoint behavior is added here.Extensibility
The sibling contract allows external teachers, model-specific alignment, and future compressed teacher representations to evolve without changing the established distillation trainer API.
Tests
Known limitations
The tensors cannot prove freshness. Reusing a prepared rollout after the student changes violates the on-policy contract, and the caller must prevent it.
Stack
Depends on the preceding upstream PR: #1745
Tunix block-diffusion design document