[maxtext] Add model-aware block-diffusion OPD#4589
Open
ethannnnnn wants to merge 4 commits into
Open
Conversation
Add block_diffusion as a default-off attention type with one namespaced block-size setting. Resolve it in the shared attention layer so ordinary global layers need no model-name dispatch while explicit specialized attention remains intact. Implement matching dense, Splash, Tokamax, packed-sequence, and load-balanced context-parallel masks. Preserve the autoregressive path and support partial final blocks. Tests: attention 42 passed/37 skipped; config 18 passed plus 21 subtests; Tokamax 3 passed; pyink, yamllint, git diff --check.
Add a model-agnostic block-diffusion training objective without changing the causal LM default. The data pipeline emits separate validity, completion, corruption, and loss masks, supports same-position/all-masked and shifted/seeded canvases, and preserves those fields through context-parallel reordering and shaped batches. Align model logits to physical targets in a shared scoring utility and consume explicit loss weights in Linen, NNX, native gradient accumulation, and the Tunix SFT adapter. The adapter composes the stock Tunix PeftTrainer and its LossOutput contract, including denominator-aware accumulation and evaluation; the existing MaxText AR trainer path remains unchanged. Tests cover partial blocks, prompt protection, deterministic corruption, alignment after sequence reordering, strict mask requirements, zero-weight gradients, causal no-regression, weighted accumulation/evaluation acceptance, and SFT adapter validation. Test Plan: - 196 passed, 14 skipped, 3 deselected; 66 subtests passed in focused MaxText unit suite - 11 passed in Tunix-backed post-training SFT suite - Pylint 10.00/10 on new scoring and loss/GA tests - Pyink, yamllint, codespell, pycompile, and git diff checks pass
Add a model-independent low-confidence block rollout that consumes target- aligned logits. It supports the two public model contracts, logical-position block boundaries after context-parallel reordering, heterogeneous batches, partial final blocks, confidence-threshold commits, and forced-argmax progress. The initial OPD scope validates a single contiguous completion suffix so clean future turns cannot leak through bidirectional intra-block attention. Shifted rollouts also require logical position zero to remain prompt context. Test Plan: - 7 passed in tests/unit/diffusion_denoise_test.py, including jax.jit execution - Pylint 10.00/10 - Pyink, pycompile, and git diff checks pass
Add a default-off student-rollout distillation source that prepares fresh block-diffusion rollouts in MaxText and delegates the weighted prepared loss to Tunix. Keep completion, corruption, validity, and loss ownership explicit; score clean generated tokens with a causal teacher; and exclude positions after model-specific stop tokens. Canonicalize multi-turn examples around the final assistant span. Earlier turns become prompt context, while any later conversation turns are truncated from tokens and both input/target segmentation before rollout. Rows without a completion span or prompt still fail closed, so the teacher and student cannot observe future context. Make resume fail closed. Persist the tokenizer, model, objective, optimizer, data-stream, topology, and stop-token contract; replay deterministic HF input before global device placement; reject early exhaustion and incompatible standard-distillation checkpoints. Honor each decoder family's native sharding mode and limit OPD to one inflight computation to control dense-logit memory. The existing dataset-driven distillation path remains the default and lazily avoids the new Tunix diffusion APIs. Test plan: - 209 passed, 14 skipped, 3 deselected, 69 subtests in the focused MaxText suite - 26 passed, 25 skipped, 22 subtests against the paired Tunix OPD head - 4 standard checkpoint restore tests passed through unittest - 45 focused MaxText OPD/input-pipeline tests and 21 subtests passed - Pyink 24.10.1, Ruff, Pylint 10/10, and git diff --check
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
ethannnnnn
requested review from
SurbhiJainUSC,
aireenmei,
darisoy,
igorts-git,
parambole,
richjames0,
shralex and
shuningjin
as code owners
July 23, 2026 22:44
ethannnnnn
requested review from
A9isha,
NuojCheng,
RissyRan,
abhinavclemson,
bvandermoon,
dipannita08,
gagika,
gobbleturk,
hengtaoguo,
huytransformer,
jiangjy1982,
khatwanimohit,
suexu1025,
vipannalla and
xibinliu
as code owners
July 23, 2026 22:44
|
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
1 participant
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 fresh student rollouts, teacher scores for exactly those generated targets, and an explicit handoff to a reusable weighted objective. Dataset-driven distillation does not provide this ownership boundary.
Scope
student_rolloutdistillation source.Design
MaxText owns rollout freshness, corruption, model-specific target alignment, teacher execution, data replay, sharding, and checkpoint identity. Tunix owns the prepared weighted KL/CE objective. Teacher logits are immutable for the loss step, and positions after configured stop tokens receive zero weight.
For multi-turn instruction data, earlier turns remain prompt context and only the final assistant span is eligible for the fresh student rollout. Token IDs after that span are cleared and both input and target segmentation mark the suffix invalid, so neither student nor teacher can observe a later user turn. Nonempty rows without both a prompt and a completion span still fail closed. Prepared batches then enter the normal MaxText placement boundary, so OPD honors the selected model's supported explicit or native auto-sharding mode instead of imposing an OPD-specific policy.
Checkpoint metadata fingerprints the student, teacher, tokenizer, objective, optimizer, data stream, topology, stop tokens, block/canvas/alignment settings, rollout policy, and loss settings. Resume replays the bounded input stream deterministically and fails on semantic mismatch or early exhaustion.
Compatibility
The existing dataset-driven distillation source remains the default and lazily avoids the new Tunix diffusion APIs. Standard distillation checkpoints and unsupported feature/evaluation combinations fail explicitly rather than being interpreted as OPD state.
Extensibility
The prepared-batch boundary can support other block-diffusion models by supplying their rollout, target alignment, and checkpoint identity. The Tunix loss remains model agnostic within its JAX/NNX execution boundary.
Tests
git diff --checkpassed.Known limitations
Dense teacher logits and one in-flight OPD computation are the correctness-first baseline. The cacheless rollout may run a full-sequence forward for each denoising transition and has not yet completed the planned real-TPU OPD performance validation. OPD intentionally trains only the final assistant span of a multi-turn example in this initial contract.
Stack
Depends on the preceding upstream PR: #4588
Cross-repository dependency: google/tunix#1747 (
block-diffusion-tunix-pr5-opdat87538fddad45).MaxText block-diffusion design document