[maxtext] Add generalized block-diffusion CFT and SFT

[ethannnnnn]

Motivation

Block-diffusion CFT/SFT needs one explicit definition of target alignment, completion scope, corruption, and loss weights. Reusing segmentation fields for these different meanings can leak clean targets or bias training.

Scope

• Add a default-off block_diffusion training objective while retaining causal_lm as the default.
• Emit distinct validity, completion, corruption, and loss masks from the data pipeline.
• Support the public same_position/all_masked and shifted/seed_and_mask model contracts.
• Consume explicit weighted targets in Linen, NNX, native gradient accumulation, and the SFT adapter.

Design

Targets remain physically unshifted. A shared scoring utility maps model logits to physical targets according to the configured alignment. Corruption samples independently per eligible block, protects prompts and padding, guarantees progress, and handles partial blocks.

completion_mask, corruption_mask, and targets_loss_mask remain separate through shaped batches and context-parallel reordering. Loss normalization uses the sum of active weights, including under gradient accumulation; an all-zero batch returns finite zero loss and zero gradients.

The Tunix-backed diffusion SFT entry point constructs the stock Tunix PeftTrainer and configures it through configure_diffusion_sft, a raw-batch adapter, and a target-aligned logits function. MaxText's existing causal-SFT compatibility subclass remains on the causal path only; diffusion does not depend on that overridden train step.

Compatibility

The causal pipeline keeps its existing trainer selection, next-token shift, segmentation-derived mask, configuration defaults, and checkpoint shape. Diffusion refuses missing explicit masks instead of guessing. Through the paired Tunix weighted-reduction dependency, the Tunix-backed path supports denominator-aware gradient accumulation and evaluation instead of averaging independently normalized microbatches.

Extensibility

The objective is selected by stable configuration rather than a model name. Additional models can implement one of the declared alignment/canvas contracts or add a new validated contract without changing trainer ownership.

Tests

• Focused MaxText suite: 196 passed, 14 skipped, 3 deselected, and 66 subtests passed.
• Tunix-backed post-training SFT suite: 11 passed.
• Tests cover partial blocks, prompt protection, deterministic corruption, reordered alignment, strict mask requirements, zero-weight gradients, and causal no-regression.
• Pylint for new scoring/loss tests, Pyink, YAML lint, codespell, Python compilation, and git diff --check passed.

Known limitations

The Tunix-backed SFT path does not support optimizer-state host offload because the underlying PeftTrainer does not provide that execution mode. Native MaxText SFT remains available for the same diffusion objective.

Stack

Depends on the preceding upstream PR: #4586

Cross-repository dependency: google/tunix#1745 (block-diffusion-tunix-pr3-sft at 8dc0dd39951a).

MaxText block-diffusion design document

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[google-cla]

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.