[maxtext] Add block-diffusion pre-training

[ethannnnnn]

Motivation

With block-diffusion primitives and block-causal attention available, MaxText
needs a focused pre-training integration before any completion fine-tuning,
rollout, distillation, replay, or RL behavior is introduced.

Dependencies

This change follows and depends on:

• [maxtext] Add block-diffusion training primitives #4737 for model-independent corruption and target alignment; and
• [maxtext] Add block-causal attention #4743 for opt-in block-causal attention.

It has no Tunix dependency.

Scope

This PR adds only:

• an opt-in training_objective=block_diffusion configuration;
• Hugging Face text pre-training corruption with separate clean targets,
  corruption masks, and target-loss masks;
• target-aligned loss handling for Linen and NNX models;
• explicit compatibility validation for the initial pre-training surface; and
• finite zero loss and gradients when an accumulated batch has no supervised
  positions.

It does not add SFT, DPO, rollout, serving, distillation, replay, or RL.

Training contract

The supported model contracts are:

• same_position logits with an all_masked canvas; and
• shifted logits with a seed_and_mask canvas.

Each nonempty eligible block keeps at least one supervised position. Clean
targets and segmentation metadata remain separate from the corrupted inputs.
Loss is normalized over targets_loss_mask positions only. Packing, SFT,
DPO, MTP, vocabulary tiling, audio, multimodal input, and non-Hugging-Face
pipelines fail validation for this initial scope.

Compatibility

training_objective=causal_lm remains the default. Its next-token data
transform, segmentation-based loss mask, and autoregressive attention path are
unchanged.

The shared gradient-accumulation helper guards a zero total-weight denominator
only for the block-diffusion objective. The causal branch retains its original
operations, including the upstream Tunix accumulation-step divisor. Focused
regression tests verify both accumulation modes against the direct full-batch
loss and gradients.

Tests

The audited current-main replay passed:

270 passed, 58 skipped, 3 deselected, 92 subtests passed

The gate covers block-diffusion primitives, dense and TPU mask construction,
configuration, Hugging Face objective selection, padding and corruption,
Linen and NNX loss masks, causal-loss fallback, shaped batches, context
parallel metadata, and positive- and zero-weight gradient accumulation.

Three pre-existing Hugging Face iterator integration tests are deselected in
the local CPU gate because they require live datasets.load_dataset access.
The new objective-transform tests in that module pass independently.

All reachable changed production statements and branches are covered by this
gate: 122/122 executable lines and 70/70 branches. The one excluded clause is
the objective-specific packing error, which cannot execute because the
block-diffusion attention validation above it necessarily rejects packing
first.

With the CI JAX/libtpu toolchain, normalized optimized HLO for DeepSeek3,
Llama3 8B, and Qwen3 1.7B is byte-identical to the public PR parent. The
feature therefore requires no reference-HLO changes.

Pyink, Pylint, Python compilation, and git diff --check are also required by
the guarded publisher.

Checklist

• I performed a self-review of the focused pre-training boundary.
• Existing autoregressive behavior remains the default and has regression coverage.
• Unsupported training modes fail closed.
• No post-training or RL implementation is included.

Design document:
https://docs.google.com/document/d/1N7KcCoAIErB2CV9EJ2G1u_mdN-AQgqwNUYSvM0mtqMI/edit

[gemini-code-assist]

Code Review

This pull request introduces support for block-diffusion pre-training within MaxText. It adds configuration options, input pipeline transformations (including token corruption and padding adjustments), and updates the pre-training loss function to handle target-aligned block-diffusion losses and explicit loss masking. It also hardens gradient accumulation against division-by-zero errors when total weights are zero, and adds comprehensive unit tests for the new functionality. There are no review comments, so I have no feedback to provide.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[hengtaoguo]

Thanks for adding the extensive test coverage. Please make sure to resolve the failing tests in CI before merging.