[tunix] Add target-aligned diffusion SFT adapter#1745
Open
ethannnnnn wants to merge 3 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.
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
Tunix needs a reusable SFT adapter for already-prepared diffusion examples without owning model-specific corruption, role parsing, or target alignment.
Scope
configure_diffusion_sftfor the canonical diffusion batch.PeftTrainerextension points.Design
The adapter validates the canonical batch, obtains target-aligned logits, sanitizes inactive targets, and returns
LossOutputwith the explicit token weights. Zero-weight positions remain numerically inert even if their placeholder values are invalid.Compatibility
Existing SFT continues to use its current data adapter, model loss, and next-token semantics unless the diffusion adapter is explicitly configured.
Extensibility
Tunix remains independent of MaxText batch keys and tokenizer policy. Any integration that can produce
DiffusionTokenBatchand target-aligned logits can use the same SFT objective.Tests
PeftTrainerwiring, JIT/gradient behavior, float32 math, no-shift alignment, inactive-value safety, and zero weights.git diff --checkpassed.Known limitations
The adapter does not create corrupted examples or infer completion scope. Those remain responsibilities of the model-aware integration.
Stack
Depends on the preceding upstream PR: #1744
Tunix block-diffusion design document