Skip to content

[tunix] Add target-aligned diffusion SFT adapter - #1891

Open
ethannnnnn wants to merge 1 commit into
google:mainfrom
ethannnnnn:block-diffusion-tunix-v2-pr5-sft
Open

[tunix] Add target-aligned diffusion SFT adapter#1891
ethannnnnn wants to merge 1 commit into
google:mainfrom
ethannnnnn:block-diffusion-tunix-v2-pr5-sft

Conversation

@ethannnnnn

Copy link
Copy Markdown
Contributor

Motivation

Tunix now has model-agnostic diffusion batch and logits contracts, but
PeftTrainer still needs a small adapter and a target-aligned supervised loss
before those contracts can be used for diffusion CFT or SFT. Autoregressive
loss helpers cannot be reused because they shift targets by one token.

Scope

This PR adds only the diffusion SFT integration layer:

  • adapt an external raw batch into DiffusionTokenBatch;
  • call a model-specific DiffusionLogitsFn through the public diffusion
    contract;
  • compute float32 cross entropy at the same token positions as target_ids;
  • preserve the unreduced weighted numerator and denominator expected by
    PeftTrainer; and
  • sanitize inactive positions so zero-weight NaN, infinity, or invalid target
    values cannot contaminate the loss or gradient.

It does not add a model implementation, tokenizer, dataset recipe,
distillation objective, rollout path, or RL trainer.

Dependencies

This PR depends on the contracts merged in #1832 and the fractional and
zero-denominator gradient behavior merged in #1833. Both dependencies are
verified against main by the guarded publisher before the patch is replayed.
It does not depend on #1850 or the optional OPD track.

Compatibility

The existing autoregressive PeftTrainer path is unchanged. Users opt in by
calling configure_diffusion_sft with their batch adapter and target-aligned
logits function. CFT and SFT use the same objective; the external adapter
chooses which target positions and weights are active.

Tests

python -m pytest -q tests/diffusion/types_test.py tests/sft/diffusion_sft_test.py
python -m pytest -q tests/sft/peft_trainer_test.py

The focused suite covers trainer wiring, same-position alignment, float32
metrics, zero-weight numerical safety and zero gradients, jitted gradients,
microbatch equivalence, and shape rejection. The guarded publisher also runs
full PeftTrainer regression tests, Pyink, scoped error/fatal Pylint, Python
compilation, and git diff --check.

Design document:
https://docs.google.com/document/d/1Xe-98ScS2RSH29AdhTdc9tO4WHQG5wCbFkJA3gmIUZQ/edit

ethannnnnn added a commit to ethannnnnn/maxtext that referenced this pull request Aug 18, 2026
Extend the opt-in block-diffusion objective from pre-training to text CFT and completion-only SFT.

Preserve clean targets and role-derived completion eligibility through the Hugging Face pipeline, then corrupt only the configured supervision scope. Reject assistant-to-user transitions within one bidirectional diffusion block to prevent future-prompt leakage.

Adapt prepared batches to the target-aligned Tunix diffusion contract, retain explicit weighted loss, treat diffusion evaluation as preaveraged, and disable NNX graph caching only for diffusion internal metrics. The causal SFT path remains the default.

Draft dependency: temporarily pin the immutable google/tunix#1891 contributor-fork head so CI can exercise the integration. Replace it with the upstream Tunix SHA before marking the PR ready.

Test Plan:
- 343 passed, 76 platform skips, 3 documented HF integration deselections, 163 subtests
- Tunix adapter has 100% statement and branch coverage
- Pyink clean and Pylint 10.00/10
- compileall, Yamllint, and git diff --check 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: diffusion contract and SFT suites; full PeftTrainer regression suite; pyink, pylint, py_compile, and diff checks.
@ethannnnnn
ethannnnnn force-pushed the block-diffusion-tunix-v2-pr5-sft branch from 0ab812a to bd8092a Compare August 19, 2026 21:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants