[maxtext] Wire opt-in block-diffusion GRPO - #4592
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
|
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. |
325c10d to
e1e18bd
Compare
| DevelopmentAndDebugging, | ||
| Profiling, | ||
| # For compatibility with trainer in post_train/rl | ||
| MaxTextConfig, |
There was a problem hiding this comment.
why all these fields removed? @hengtaoguo PTAL.
There was a problem hiding this comment.
Based on the updated doc string below, it says inheriting MaxTextConfig and RL-specific mixins will cover all the above deleted args, so the pre/post-training won't be affected. Looks reasonable based on the design, but I would recommend author to double check whether the config was impacted or not.
|
|
||
|
|
||
| class TrainingLoop(BaseModel): | ||
| class BlockDiffusionObjective(BaseModel): |
There was a problem hiding this comment.
@hengtaoguo can you help me review these changes.
There was a problem hiding this comment.
Great catch, I would prefer not to change the TrainingLoop from inheriting BaseModel to BlockDiffusionObjective.
I understand this consolidated design choice, but introducing a nested hierarchy only for block-diffusion sets an unusual example. I would recommend author to either flatten these fields directly into TrainingLoop, or use sibling mixin. It keeps the configuration aligned with the current style. You could add a comment line to indicate such fields are for block diffusion purpose.
| completion_only=sft_train_on_completion_only, | ||
| max_target_length=max_target_length, | ||
| unk_id=pad_id, | ||
| target_aligned=is_block_diffusion, |
There was a problem hiding this comment.
SFTPromptMasking should accept the training_objective string directly. The transform itself should internally determine if it needs to align targets or emit completion masks based on the objective, keeping the pipeline construction clean.
| raise ValueError(f"{name} must match initial_tokens shape; received {tuple(value.shape)} and {expected_shape}") | ||
|
|
||
|
|
||
| def _concrete_numpy(value): |
There was a problem hiding this comment.
create a utility file and add _concrete_numpy as it is duplicated multiple times in the code.
| @@ -0,0 +1,348 @@ | |||
| # Copyright 2026 Google LLC | |||
There was a problem hiding this comment.
the class name is to ambiguous. Since diffusion is broad (continuous/discrete), add a subfolder for block diffusion and add the block diffusion specific classes in there. Rename it to something like BlockDiffusionRolloutTrace or more broad DiscreteUnmaskTrace.
| @@ -0,0 +1,43 @@ | |||
| # Copyright 2026 Google LLC | |||
There was a problem hiding this comment.
like denoise.py, the class name is very broad, add it under a block_diffusion folder and rename it to something like block_alignment or target_alignment.
| return jnp.sum(batch.get("targets_loss_mask", batch["targets_segmentation"]) != 0) | ||
|
|
||
| @override | ||
| def eval_loss_is_preaveraged(self) -> bool: |
There was a problem hiding this comment.
loss preaveraging should be decoupled from block_diffusion, instead add a new config flag loss_is_preaveraged so that it can be used more broadly.
| mask_kwargs = {"cp_size": cp_size} if use_load_balanced_cp else {} | ||
| mask = mask_type( | ||
| shape=mask_shape, | ||
| block_size=self.block_diffusion_block_size, |
There was a problem hiding this comment.
rename this to causal_block_size to make it more generic.
e1e18bd to
8adfa6e
Compare
|
Addressed the current review feedback and rebased the complete seven-PR stack onto the latest Key changes in this update:
Validation: 209 focused tests passed, 44 skipped, with 71 parameterized subtests; cumulative Pylint scored 10.00/10; Pyink, compilation, whitespace, and public-text checks passed. This PR retains its documented dependency on Tunix PR #1745. Design: https://docs.google.com/document/d/1N7KcCoAIErB2CV9EJ2G1u_mdN-AQgqwNUYSvM0mtqMI/edit Vanilla public reproduction guide: https://docs.google.com/document/d/1rjRSopy888ZqLnGh_KUseW8h1_ZqPGEuSee5MM2Q9ac/edit @entrpn @hengtaoguo, could you please take another look? |
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
Extend the low-confidence block denoiser with a stochastic, per-row rollout that records the sampled token, action step, and action log probability for every completion position. The trace keeps the full fixed-horizon denoising schedule even after a visible EOS, isolates RNG streams between batch rows, excludes the mask token from sampling, and forces progress when no token clears the confidence threshold. Existing deterministic generation remains unchanged. This trace is the minimal model-independent contract needed to replay the same diffusion actions under live, reference, and old policies during policy optimization. Test Plan: - pytest tests/unit/diffusion_denoise_test.py - included in the 83-test focused diffusion-RL suite - Pyink, Pylint, py_compile, and git diff --check
Add a correctness-first MaxText/Tunix adapter for block-diffusion policy optimization. The in-process rollout exports the full stochastic denoising trace while keeping user-visible completions truncated at the first stop token. The scorer reconstructs each pre-commit canvas and evaluates the recorded action at its original denoising step, so rollout, live, reference, and old-policy log probabilities share one target-aligned contract. Generation uses independent train/eval RNG streams with resumable step contexts. Invalid stop IDs, non-finite logits, unresolved masks, unsupported sampling filters, and partial/non-addressable host traces fail closed. The adapter is opt-in and does not register or import the vLLM path. Test Plan: - pytest tests/post_training/unit/diffusion_rl_test.py - included in the 83-test focused diffusion-RL suite - Pyink and changed-file Pylint 10.00/10 - py_compile and git diff --check
Expose the exact denoising-trace rollout and scorer through MaxText's existing Tunix RL workflow. Add explicit diffusion-RL configuration with default-off validation, construct the custom rollout only for block-diffusion objectives, and pass its prepared scorer to the GRPO learner. Train and evaluation generation stay inside the RLCluster mode contract, use separate sampling configurations, and avoid any vLLM import or registration on the diffusion path. Checkpoint metadata fingerprints every setting that changes rollout or replay semantics. Resume therefore fails closed across model, tokenizer, mask, block, alignment, canvas, temperature, sampling, stop-token, and trace-schedule changes. Autoregressive defaults and the existing vLLM workflow remain unchanged. Initial MaxText routing supports GRPO and token-level GSPO. Other Tunix GRPO family variants can consume the prepared scorer but are not exposed here yet. Test Plan: - focused MaxText diffusion-RL suite: 83 passed, 3 subtests passed - Pyink and changed-file Pylint 10.00/10 - py_compile and git diff --check
8adfa6e to
b28b35c
Compare
3 participants
Motivation
The exact rollout and replay contracts need to enter the existing MaxText/Tunix RL workflow without changing autoregressive GRPO or comparing incompatible probability semantics.
Scope
RLClustermode contract.Design
MaxText supplies a custom rollout model plus
diffusion_logits_fnto Tunix. Rollout, live actor, detached reference actor, and the detached start-of-step actor snapshot used as the old policy score one full target-aligned trace. The snapshot remains on the actor role; it is not a separate model role. Training and evaluation have independent sampling configurations and generation contexts tied to global step, so resume does not silently change RNG progression.Checkpoint metadata covers model, tokenizer, mask token, block size, alignment, canvas policy, temperature, filters, threshold, maximum denoising steps, stop IDs, sequence lengths, and full-trace semantics.
RLConfigderives from the completeMaxTextConfigmodel schema and adds only the RL-specific configuration groups. This keeps RL parsing aligned as model fields evolve, without maintaining a partial duplicate list or admitting arbitrary unknown fields.Block-diffusion CFT/SFT and OPD still require the Hugging Face MaxText input pipeline. The dedicated diffusion-RL rollout uses the RL dataset loader, so it alone is exempt from
dataset_type=hf;packing=Falseremains mandatory for every block-diffusion objective.Tunix reshards model variables onto the sampler role mesh, while MaxText modules also retain graph-static sharding objects. The rollout constructor recursively rebinds direct
Meshvalues,MeshorNamedShardingvalues cached infunctools.partialarguments and keywords, and precomputedNamedShardingfields before compilation. It logs the number of rebound module fields. This changes neither parameter values nor partition specs and does not hardcode a model or device ID.Compatibility
Autoregressive configuration, vLLM registration, and existing rollout behavior remain unchanged when
rl.diffusion_rolloutis false. Unsupported agentic, packing, MTP, sampling-filter, or execution-mode combinations fail before training.Full Qwen3 model configuration is validated through the same typed schema used by MaxText model construction; no model-name exception or permissive extra-field mode is introduced. Autoregressive RL still uses the existing configuration and rollout defaults.
Graph-mesh rebinding is confined to the custom in-process diffusion rollout. Existing vLLM and autoregressive rollout construction is unchanged.
Extensibility
The initial MaxText routing supports GRPO and token-level GSPO. Tunix owns algorithm variants, advantage calculation, and optimizer orchestration; MaxText owns model-aware trace preparation and checkpoint identity.
Tests
diffusion_rl_test.pysuite with the fix: 14 passed on two CPU devices.Meshbinder deterministically reproduced JAX's incompatible-device error: the direct mesh and model state were on CPU device ID 1 while a cached partial'sNamedShardingstill referenced CPU device ID 0.RLConfigvalidation: 14 passed, including the HF-only exemption for the dedicated diffusion-RL loader while preserving the packing rejection.RLConfigwithemb_dim=1024, 28 layers,use_mrope=false, anduse_qk_norm=true.RLConfig.six optional vLLM/SGL-JAX integration tests deselected. A 4x4x4 TPU
validation then completed on all 16 workers with
zero restarts and all three configured actor updates. Losses and rewards
were finite, importance-sampling ratio was 1, clip fraction was 0, and
gradient norms were nonzero. The actor step-3 checkpoint finalized and the
downstream model export was verified. This is correctness smoke evidence,
not a throughput or model-quality benchmark.
git diff --checkpassed for the sampler-mesh change; the earlier changed-file Pylint, compilation, and YAML checks remain green.Known limitations
MaxText does not yet expose DAPO or DrGRPO routing for this path. The current implementation requires Pathways single-controller execution and does not support packing or agentic rollout. Recursive graph-static sharding rebinding is covered locally and the three-update TPU smoke completed with checkpoint/export verification, but broader throughput, memory scaling, convergence, and model-quality validation remain open.
Stack
Depends on the preceding upstream PR: #4591
Cross-repository dependency: google/tunix#1749 (
block-diffusion-tunix-pr7-policy-optimizationat264e9f88b374).MaxText block-diffusion design document
Checklist