[maxtext] Add block-diffusion CFT and SFT - #4891
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request integrates target-aligned block-diffusion Supervised Fine-Tuning (SFT) with Tunix. It updates configuration types, data processing pipelines, and training hooks to support block-diffusion SFT, including handling of completion-only masking and preaveraged evaluation losses. Additionally, a new adapter module diffusion_sft.py is introduced to interface MaxText with Tunix's diffusion SFT. Feedback on the changes highlights a performance concern in the newly added adapter, where eager validation via np.asarray on JAX arrays could trigger device-to-host transfers and degrade training throughput.
| def _concrete_numpy(value): | ||
| """Returns a host view only when eager validation is safe.""" | ||
| if isinstance(value, jax.core.Tracer): | ||
| return None | ||
| if isinstance(value, jax.Array) and not value.is_fully_addressable: | ||
| return None | ||
| return np.asarray(value) |
There was a problem hiding this comment.
Calling np.asarray on JAX arrays inside _concrete_numpy on every training step triggers device-to-host transfers and host-device synchronization, which can severely degrade training throughput. Consider disabling these eager validations by default in production, or only running them on the first step, or gating them behind a debug configuration flag.
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
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
6957f93 to
cd571e1
Compare
2 participants
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:
Description
Start with a short description of what the PR does and how this is a change from
the past.
The rest of the description includes relevant details and context, examples:
If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456
You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/123456
Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.
Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.
Tests
Please describe how you tested this change, and include any instructions and/or
commands to reproduce.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.