diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py index 57b592f026..94fd67dba4 100644 --- a/src/maxtext/common/common_types.py +++ b/src/maxtext/common/common_types.py @@ -123,6 +123,7 @@ class AttentionType(enum.Enum): MLA = "mla" COMPRESSED = "compressed" FULL = "full" + BLOCK_DIFFUSION = "block_diffusion" class ShardMode(enum.Enum): diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 9c834fb7d4..d96b0aab9c 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -374,12 +374,21 @@ param_scan_axis: 1 # The attention parameter dictates the specific algorithm/methodology used to compute the attention scores # The attention_type parameter determines the variants of attention, e.g. global or local_sliding attention: 'autoselected' # Supported attention: autoselected, dot_product, flash, cudnn_flash_te -attention_type: 'global' # Supported attention_type: global, local_sliding, chunk, mla +attention_type: 'global' # Supported attention_type: global, local_sliding, chunk, mla, full, compressed, block_diffusion share_kv_projections: false # Note: Not compatible with attention_type='mla' attention_bias: false # If true, adds a learnable bias to the query, key, and value projections attention_sink: false sliding_window_size: 0 chunk_attn_window_size: 0 +# Block-diffusion attention (Fast_dLLM / Diffusion-GR2): bidirectional within a block of bd_size tokens, block-causal across blocks. Weight-preserving. +enable_block_diffusion: false +bd_size: 32 # Block size for block-diffusion attention (must be > 0 when enabled; max_target_length must be divisible by it) +mask_id: 151669 # Token id of the [MASK] token used by block-diffusion (must be < vocab_size when enabled) +# On-policy distillation (OPD): the block-diffusion student generates its own completion by denoising, the frozen teacher scores it, and forward-KL is applied at the committed positions. Requires enable_block_diffusion on the student. +opd_on_policy: false +opd_threshold: 0.9 # Confidence threshold for committing a denoised position during the on-policy rollout. +opd_temperature: 0.0 # Softmax temperature for the on-policy commit step (0.0 = greedy/argmax). +opd_max_denoise_steps: 0 # Per-block denoise iteration cap for the on-policy rollout; <= 0 means use bd_size. attn_logits_soft_cap: 0.0 final_logits_soft_cap: 0.0 z_loss_multiplier: 0.0 diff --git a/src/maxtext/configs/post_train/rl.yml b/src/maxtext/configs/post_train/rl.yml index 64460ea066..ee1de57b94 100644 --- a/src/maxtext/configs/post_train/rl.yml +++ b/src/maxtext/configs/post_train/rl.yml @@ -72,6 +72,17 @@ rl: # If null, the entire model is resharded at once. reshard_chunk_size: null + # ====== Block-diffusion RL rollout (Fast_dLLM / Diffusion-GR2, TraceRL-style) ====== + # If true, the policy is a block-diffusion model: rollouts generate by block-diffusion + # denoising (MaxTextDiffusionRollout) and the GRPO/DAPO importance ratio uses the shared + # block-diffusion per-token logprob instead of the AR next-token logprob. Requires + # enable_block_diffusion=true (top level). When false, RL is autoregressive and unchanged. + rl_diffusion_rollout: false + # Confidence threshold for committing a denoised position during the diffusion rollout. + rl_diffusion_threshold: 0.9 + # Per-block denoise iteration cap for the diffusion rollout; <= 0 means use bd_size. + rl_diffusion_max_denoise_steps: 0 + # ====== Models ====== # for MaxText diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 3dffe5e7ac..37fb47b2a0 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -559,7 +559,7 @@ class Attention(BaseModel): "autoselected", description="The attention algorithm to use (dot_product, flash, etc).", ) - attention_type: Literal["global", "local_sliding", "chunk", "mla", "full", "compressed"] = Field( + attention_type: Literal["global", "local_sliding", "chunk", "mla", "full", "compressed", "block_diffusion"] = Field( "global", description="The variant of attention to use." ) share_kv_projections: bool = Field( @@ -606,6 +606,17 @@ class Attention(BaseModel): ) sliding_window_size: NonNegativeInt = Field(0, description="The size of the sliding window for local attention.") chunk_attn_window_size: NonNegativeInt = Field(0, description="The window size for chunked attention.") + enable_block_diffusion: bool = Field( + False, + description=( + "If True, use block-diffusion attention: bidirectional within a block of `bd_size` tokens " + "and block-causal across blocks (Fast_dLLM / Diffusion-GR2). Weight-preserving." + ), + ) + bd_size: int = Field(32, description="Block size for block-diffusion attention. Must be > 0 when enabled.") + mask_id: int = Field( + 151669, description="Token id of the [MASK] token used by block-diffusion. Must be < vocab_size when enabled." + ) attn_logits_soft_cap: None | NonNegativeFloat = Field( None, description="Soft-cap value for attention logits. None means no cap." ) @@ -1478,6 +1489,26 @@ class Distillation(BaseModel): "The other parameters will be frozen if this attribute is non empty)", ) + # --- On-policy distillation (OPD) for block-diffusion -- + opd_on_policy: bool = Field( + False, + description=( + "If True, run on-policy block-diffusion distillation: the student generates its own completion " + "by block-diffusion denoising, the frozen teacher scores that sequence, and forward-KL is applied " + "only at the committed positions. Requires enable_block_diffusion on the student. When False, " + "distillation is unchanged (off-policy)." + ), + ) + opd_threshold: float = Field( + 0.9, description="Confidence threshold for committing a denoised position during the on-policy rollout." + ) + opd_temperature: float = Field( + 0.0, description="Softmax temperature for the on-policy commit step (0.0 = greedy/argmax)." + ) + opd_max_denoise_steps: int = Field( + 0, description="Per-block denoise iteration cap for the on-policy rollout; <= 0 means use bd_size." + ) + class TrainingLoop(BaseModel): """Configuration for the main training loop, evaluation, and reproducibility.""" @@ -2173,6 +2204,21 @@ class RL(BaseModel): "If None, no chunking is applied, which may lead to OOM errors if tensors are too large." ), ) + rl_diffusion_rollout: bool = Field( + False, + description=( + "If True, the policy is a block-diffusion model: rollouts generate by block-diffusion " + "denoising (MaxTextDiffusionRollout) and the GRPO/DAPO importance ratio uses the shared " + "block-diffusion (TraceRL-style) per-token logprob instead of the AR next-token logprob. " + "Requires enable_block_diffusion=True. When False, RL is autoregressive and unchanged." + ), + ) + rl_diffusion_threshold: float = Field( + 0.9, description="Confidence threshold for committing a denoised position during the RL diffusion rollout." + ) + rl_diffusion_max_denoise_steps: int = Field( + 0, description="Per-block denoise iteration cap for the RL diffusion rollout; <= 0 means use bd_size." + ) class RLDataset(BaseModel): @@ -3510,3 +3556,28 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de f"must be equal to attention_output_dim ({self.attention_output_dim})" ) return self + + @model_validator(mode="after") + def _validate_block_diffusion(self) -> "MaxTextConfig": + """Validates block-diffusion attention settings. + + Only enforced when `enable_block_diffusion` is True, because the block-diffusion + defaults (e.g. `mask_id`) are intentionally out of range for the default vocab. + Defined after `set_derived_and_validate_values`; `vocab_size` and + `max_target_length` are never mutated by that validator, so their final values + are visible here regardless of validator ordering. + """ + if self.enable_block_diffusion: + if not isinstance(self.bd_size, int) or self.bd_size <= 0: + raise ValueError("`bd_size` must be an integer > 0 when `enable_block_diffusion` is True.") + if self.mask_id >= self.vocab_size: + raise ValueError( + f"`mask_id` ({self.mask_id}) must be < `vocab_size` ({self.vocab_size}) " + "when `enable_block_diffusion` is True." + ) + if self.max_target_length % self.bd_size != 0: + raise ValueError( + f"`max_target_length` ({self.max_target_length}) must be divisible by `bd_size` " + f"({self.bd_size}) when `enable_block_diffusion` is True." + ) + return self diff --git a/src/maxtext/input_pipeline/hf_data_processing.py b/src/maxtext/input_pipeline/hf_data_processing.py index 370f1895bd..5895ec054a 100644 --- a/src/maxtext/input_pipeline/hf_data_processing.py +++ b/src/maxtext/input_pipeline/hf_data_processing.py @@ -357,7 +357,19 @@ def preprocessing_pipeline( operations.append(input_pipeline_utils.PadOrTrimToMaxLength(max_target_length, pad_id)) operations.append(grain.Batch(batch_size=batch_size, drop_remainder=drop_remainder)) - if shift and not use_dpo: + if config.enable_block_diffusion and not use_dpo: + # Masked-diffusion (CFT) corruption for block-diffusion SFT (Fast_dLLM / Diffusion-GR2). + # Replaces ShiftData: block-diffusion targets are aligned, not next-token shifted. + if packing: + raise ValueError("enable_block_diffusion is not supported with packing=True; set packing=False.") + operations.append( + input_pipeline_utils.BlockDiffusionMasking( + bd_size=config.bd_size, + mask_id=config.mask_id, + axis=1, + ) + ) + elif shift and not use_dpo: operations.append(input_pipeline_utils.ShiftData(ignored_ids=[pad_id, tokenizer.bos_token_id], axis=1)) # Since HuggingFace IterableDataset does not support access through index diff --git a/src/maxtext/input_pipeline/input_pipeline_utils.py b/src/maxtext/input_pipeline/input_pipeline_utils.py index 8ba57363fb..787edcb666 100644 --- a/src/maxtext/input_pipeline/input_pipeline_utils.py +++ b/src/maxtext/input_pipeline/input_pipeline_utils.py @@ -976,6 +976,95 @@ def map(self, element): return shift_and_refine(element, ignored_ids=self.ignored_ids, axis=self.axis) +@dataclasses.dataclass +class BlockDiffusionMasking(grain.RandomMapTransform): + """Masked-diffusion (conversion fine-tuning / CFT) corruption for block-diffusion SFT. + + Turns an autoregressive SFT batch into a masked-diffusion training example + (Fast_dLLM / Diffusion-GR2). For each example a random subset of the RESPONSE + tokens is corrupted by replacing them with ``mask_id`` in ``inputs``; the clean + tokens are kept as ``targets`` ALIGNED (NOT next-token shifted); and + ``targets_segmentation`` is set to 1 only at the masked positions. MaxText's loss + already gates the per-position cross-entropy on ``targets_segmentation != 0``, so + writing the "predict-here" mask into ``targets_segmentation`` makes the loss + compute exactly (and only) at the masked positions with NO loss-side change. + ``inputs_segmentation`` (the attention mask) is left untouched so attention still + sees every real prompt+response token, including the masked ones being denoised. + + This transform REPLACES ``ShiftData`` in the block-diffusion pipeline: block + diffusion predicts each masked position from the (bidirectional-within-block) + context at the SAME position, so targets must stay aligned rather than shifted. + + Response span: the maskable span is exactly the completion tokens, identified as + ``targets_segmentation != 0`` on the incoming (post ``PadOrTrimToMaxLength``, + pre-shift) batch. With ``sft_train_on_completion_only=True`` this is the + assistant/completion span (prompt and padding have ``targets_segmentation == 0`` + and are never masked); with completion-only disabled (or non-SFT) it is every + real token, matching a LLaDA-style whole-sequence masked-diffusion objective. + + Masking scheme (v1, per-block; matches the Fast_dLLM reference): the sequence is + split into contiguous blocks of ``bd_size`` tokens, aligned to position 0 so the + blocks coincide with the block-diffusion attention blocks (``position // bd_size``). + For each block a ratio ``t ~ U(0, 1)`` is drawn and the per-token mask probability + is ``p = (1 - eps) * t + eps``. The ``eps`` floor (1e-3) keeps ``p`` strictly + positive so a block is never guaranteed fully unmasked. Every RESPONSE position in + the block is then masked independently with probability ``p``. Drawing ``t`` + per-block (rather than one ratio for the whole sequence) gives a mix of lightly and + heavily corrupted blocks within a single example, which is what the block-diffusion + decoder sees at inference (block-by-block denoising). + + Randomness uses grain's per-record RNG (``grain.RandomMapTransform``) so each batch + draws a fresh, reproducible mask; masks for the examples within a batch are drawn + independently. For a fixed ``rng`` the output is deterministic. + + TODO(block-diffusion): add the optional complementary-mask pass from the Fast_dLLM + reference (also train on the un-masked response tokens using the complement of the + sampled mask) for a lower-variance objective. Skipped in v1. + """ + + def __init__(self, bd_size, mask_id, eps=1e-3, axis=1): + if bd_size <= 0: + raise ValueError(f"bd_size must be positive, got {bd_size}") + self.bd_size = bd_size + self.mask_id = mask_id + self.eps = eps + self.axis = axis + + def random_map(self, element, rng: np.random.Generator): + inputs = element["inputs"] + targets = element["targets"] + targets_segmentation = element["targets_segmentation"] + + seq_len = inputs.shape[self.axis] + if seq_len % self.bd_size != 0: + raise ValueError(f"sequence length ({seq_len}) must be divisible by bd_size ({self.bd_size})") + num_blocks = seq_len // self.bd_size + + # Maskable span = completion tokens (targets_segmentation != 0). Prompt and + # padding carry targets_segmentation == 0 and are therefore never masked. + response = targets_segmentation != 0 + + # Per-block ratio t ~ U(0,1) -> per-token mask probability p = (1-eps)*t + eps, + # broadcast from one value per block to every position in that block. + block_shape = list(inputs.shape) + block_shape[self.axis] = num_blocks + t = rng.random(size=tuple(block_shape)) + p_block = (1.0 - self.eps) * t + self.eps + p_pos = np.repeat(p_block, self.bd_size, axis=self.axis) + + mask = (rng.random(size=inputs.shape) < p_pos) & response + + element["inputs"] = np.where(mask, inputs.dtype.type(self.mask_id), inputs) + # targets stay aligned to inputs (clean tokens, no shift). Values outside the + # masked positions are ignored by the loss (targets_segmentation == 0 there). + element["targets"] = targets + # Loss is computed exactly at the masked positions. + element["targets_segmentation"] = mask.astype(targets_segmentation.dtype) + # inputs_segmentation / *_position are intentionally left unchanged: attention + # must still cover every real prompt+response token, including masked ones. + return element + + @dataclasses.dataclass class ComputeQwen3OmniPositions(grain.MapTransform): """Computes 3D position IDs for Qwen3-Omni multimodal sequences. diff --git a/src/maxtext/integration/vllm/maxtext_diffusion_rollout.py b/src/maxtext/integration/vllm/maxtext_diffusion_rollout.py new file mode 100644 index 0000000000..b993398d05 --- /dev/null +++ b/src/maxtext/integration/vllm/maxtext_diffusion_rollout.py @@ -0,0 +1,324 @@ +# Copyright 2023–2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Block-diffusion RL rollout for GRPO/DAPO (TraceRL-style). + +``MaxTextVllmRollout`` (in ``maxtext_vllm_rollout.py``) drives autoregressive +generation through a KV-cached vLLM engine and returns the vLLM sampler's +next-token logprobs. A block-diffusion policy is NOT autoregressive: it generates +a completion by *denoising* a masked canvas block-by-block, and its per-token +log-probability must come from a bidirectional-within-block forward, not an AR +decode. This module therefore provides a fresh JAX rollout (a plain +``tunix.rl.rollout.base_rollout.BaseRollout``, NOT a ``VllmRollout`` subclass): + + * ``generate`` runs block-diffusion denoising via + ``diffusion_generate.diffusion_generate``; + * ``get_per_token_logps`` returns the shared block-diffusion logprob + (``diffusion_generate.diffusion_per_token_logps``). + +THE P0 INVARIANT. The GRPO importance ratio ``exp(new_logp - old_logp)`` is only +unbiased if the *old* logp (this rollout's ``get_per_token_logps``) and the *new* +logp (computed in the tunix loss) come from the SAME function. Both funnel through +``diffusion_per_token_logps`` via ``_diffusion_per_token_logps_from_model`` below: +this rollout calls it directly, and the tunix loss calls it through the pluggable +hook built by ``make_diffusion_per_token_logps_fn`` (installed on the GRPO learner +when the diffusion rollout is active). tunix never imports maxtext — the shared fn +lives here and is injected into tunix as a callback. +""" + +from __future__ import annotations + +import functools +from typing import Any, Optional, Tuple + +import jax +import jax.numpy as jnp +import numpy as np +from flax import nnx + +from tunix.rl.rollout import base_rollout + +from maxtext.trainers.post_train.distillation import diffusion_generate + + +def _diffusion_model_apply(model: Any) -> diffusion_generate.ForwardFn: + """Builds the cacheless block-diffusion forward ``(tokens, positions) -> logits``. + + ``model`` is the tunix-wrapped MaxText actor (``TunixMaxTextAdapter``), whose + call signature is ``(input_tokens, positions, cache, attention_mask) -> (logits, + None)``. The adapter ignores ``attention_mask`` (MaxText applies its own + block-diffusion mask internally when ``enable_block_diffusion`` is on) and, given + ``decoder_segment_ids=None``, synthesizes pad-masking segment ids from its + ``pad_id``. Both RL sites (rollout + loss hook) build ``model_apply`` this way, so + the forward is identical on both sides of the importance ratio. + """ + + def model_apply(tokens: jax.Array, positions: jax.Array) -> jax.Array: + out = model(tokens, positions, None, None) + return out[0] if isinstance(out, tuple) else out + + return model_apply + + +def _diffusion_per_token_logps_from_model( + model: Any, + prompt_tokens: jax.Array, + completion_tokens: jax.Array, + completion_mask: jax.Array | None = None, + *, + bd_size: int, + pad_id: int, + temperature: float = 1.0, + stop_gradient: bool = False, +) -> jax.Array: + """The one bridge from a (wrapped) MaxText model to the shared logprob fn. + + Both P0 sites call THIS helper so they are byte-for-byte the same code path into + ``diffusion_generate.diffusion_per_token_logps`` (the single source of truth). + """ + return diffusion_generate.diffusion_per_token_logps( + _diffusion_model_apply(model), + prompt_tokens, + completion_tokens, + completion_mask, + bd_size=bd_size, + pad_id=pad_id, + temperature=temperature, + stop_gradient=stop_gradient, + ) + + +def make_diffusion_per_token_logps_fn(bd_size: int, mask_id: int | None = None): + """Builds the tunix ``per_token_logps_fn`` hook for a block-diffusion policy. + + The returned callable is a drop-in for ``tunix.rl.common.compute_per_token_logps`` + (same call signature, called with ``graphdef, state, ...`` from ``grpo_loss_fn``), + but computes the block-diffusion logprob instead of the AR one. It reconstructs + the model from ``graphdef``/``state`` and funnels into + ``_diffusion_per_token_logps_from_model`` — the exact path the rollout uses — so + the new-policy and old-policy logps are guaranteed identical (the P0 invariant). + + ``mask_id`` is accepted for symmetry with the rollout config; it is only needed + for generation (seeding the canvas), not for scoring an already-committed + sequence, so it is unused here. + """ + del mask_id # only used by generate(); scoring needs no mask token + + @functools.partial( + jax.jit, + static_argnames=( + "pad_id", + "eos_id", + "stop_gradient", + "return_logits", + "temperature", + ), + ) + def diffusion_per_token_logps_hook( + graphdef, + state, + prompt_tokens: jax.Array, + completion_tokens: jax.Array, + pad_id: int, + eos_id: int, + images: jax.Array | None = None, + completion_mask: jax.Array | None = None, + stop_gradient: bool = True, + return_logits: bool = False, + segment_ids: jax.Array | None = None, + segment_positions: jax.Array | None = None, + temperature: float = 1.0, + ): + # Block-diffusion RL scores the plain prompt|completion sequence; packing + # (segment_ids), images and eos-based masking are AR-path concerns. + del images, eos_id, segment_ids, segment_positions + model = nnx.merge(graphdef, state) + logps = _diffusion_per_token_logps_from_model( + model, + prompt_tokens, + completion_tokens, + completion_mask, + bd_size=bd_size, + pad_id=pad_id, + temperature=temperature, + stop_gradient=stop_gradient, + ) + if return_logits: + return logps, None + return logps + + return diffusion_per_token_logps_hook + + +class MaxTextDiffusionRollout(base_rollout.BaseRollout): + """Block-diffusion RL rollout (denoising generation + shared diffusion logprob). + + Constructed by ``RLCluster`` via the same custom-class path as + ``MaxTextVllmRollout`` — i.e. ``rollout_engine=functools.partial( + MaxTextDiffusionRollout, maxtext_config=trainer_config)`` — which calls it with + ``rollout_actor=``, ``tokenizer=``, ``mesh=``, ``rollout_config=``. + """ + + def __init__( + self, + rollout_actor: Any, + tokenizer: Any, + mesh: jax.sharding.Mesh, + rollout_config: base_rollout.RolloutConfig, + maxtext_config: Any, + cache_config_or_size: base_rollout.CacheConfig | int | None = None, + ): # pylint: disable=too-many-positional-arguments + del cache_config_or_size # block-diffusion rollout is cacheless + self._model = rollout_actor + self._tokenizer = tokenizer + self._mesh = mesh + self._rollout_config = rollout_config + self._config = maxtext_config + + if not getattr(maxtext_config, "enable_block_diffusion", False): + raise ValueError( + "MaxTextDiffusionRollout requires enable_block_diffusion=True on the " + "policy model (the rollout denoises a block-diffusion canvas)." + ) + # bd_size / mask_id / enable_block_diffusion are top-level (flattened + # Attention config); the diffusion-rollout generation knobs live in the nested + # `rl` config, so read them off that sub-config. + self._bd_size = maxtext_config.bd_size + self._mask_id = maxtext_config.mask_id + _rl = getattr(maxtext_config, "rl", None) + self._threshold = getattr(_rl, "rl_diffusion_threshold", 0.9) if _rl is not None else 0.9 + self._max_denoise_steps = getattr(_rl, "rl_diffusion_max_denoise_steps", 0) if _rl is not None else 0 + self._temperature = rollout_config.temperature + self._pad_id = tokenizer.pad_id() + self._eos_id = tokenizer.eos_id() + + @property + def mesh(self) -> jax.sharding.Mesh: + return self._mesh + + def _encode_left_padded(self, prompts: list[str], max_prompt_length: int) -> np.ndarray: + """Tokenizes and left-pads/truncates each prompt to ``[B, max_prompt_length]``.""" + rows = [] + for prompt in prompts: + ids = list(self._tokenizer.encode(prompt)) + ids = ids[-max_prompt_length:] # truncate on the left (keep the most recent) + pad = [self._pad_id] * (max_prompt_length - len(ids)) + rows.append(pad + ids) + return np.asarray(rows, dtype=np.int32) + + def generate( + self, + prompts: list[str], + rollout_config: base_rollout.RolloutConfig, + **kwargs, + ) -> base_rollout.RolloutOutput: + """Generates completions by block-diffusion denoising.""" + del kwargs + max_prompt_length = rollout_config.max_prompt_length + max_gen = rollout_config.max_tokens_to_generate + + prompt_ids = self._encode_left_padded(prompts, max_prompt_length) # [B, P] + batch_size = prompt_ids.shape[0] + + # Pad the canvas so prompt+completion is a whole number of blocks (the + # block-diffusion / diffusion_generate contract). Completion tokens beyond + # max_gen are dropped from the returned output. + total = max_prompt_length + max_gen + padded_total = -(-total // self._bd_size) * self._bd_size # round up to bd_size + canvas_len = padded_total + completion_span = canvas_len - max_prompt_length + + # Seed canvas: prompt (left-padded) then mask_id over the completion span. + canvas = np.full((batch_size, canvas_len), self._mask_id, dtype=np.int32) + canvas[:, :max_prompt_length] = prompt_ids + canvas = jnp.asarray(canvas) + + # Generation validity: real prompt tokens + the whole completion span. + prompt_valid = jnp.asarray(prompt_ids != self._pad_id) + gen_mask = jnp.concatenate( + [ + jnp.zeros((batch_size, max_prompt_length), dtype=bool), + jnp.ones((batch_size, completion_span), dtype=bool), + ], + axis=1, + ) + attend_mask = jnp.concatenate([prompt_valid, jnp.ones((batch_size, completion_span), dtype=bool)], axis=1) + positions = diffusion_generate._positions_from_mask(attend_mask) # pylint: disable=protected-access + + gen_tokens, _ = diffusion_generate.diffusion_generate( + _diffusion_model_apply(self._model), + canvas, + positions, + gen_mask, + mask_id=self._mask_id, + bd_size=self._bd_size, + threshold=self._threshold, + temperature=self._temperature, + max_denoise_steps=self._max_denoise_steps, + ) + completions = np.asarray(gen_tokens[:, max_prompt_length : max_prompt_length + max_gen]) + + # Truncate each completion at its first EOS (inclusive) and decode. + tokens_out: list[np.ndarray] = [] + text_out: list[str] = [] + for row in completions: + row = np.asarray(row, dtype=np.int32) + eos_hits = np.flatnonzero(row == self._eos_id) + end = int(eos_hits[0]) + 1 if eos_hits.size > 0 else row.shape[0] + trimmed = row[:end] + tokens_out.append(trimmed) + text_out.append(self._tokenizer.decode(trimmed.tolist())) + + return base_rollout.RolloutOutput( + text=text_out, + logits=None, + tokens=tokens_out, + left_padded_prompt_tokens=prompt_ids, + logprobs=None, + ) + + def get_per_token_logps( + self, + prompt_tokens: jax.Array, + completion_tokens: jax.Array, + completion_mask: jax.Array | None = None, + ) -> jax.Array: + """Old-policy block-diffusion logps (P0 site #1 — shared with the loss).""" + return _diffusion_per_token_logps_from_model( + self._model, + prompt_tokens, + completion_tokens, + completion_mask, + bd_size=self._bd_size, + pad_id=self._pad_id, + temperature=self._temperature, + stop_gradient=True, + ) + + def update_params( + self, + params: Any, + filter_types: Optional[Tuple[Any, ...]] = None, + ) -> None: + del filter_types + nnx.update(self._model, params) + + def pad_id(self) -> int: + return self._pad_id + + def eos_id(self) -> int: + return self._eos_id + + def model(self) -> nnx.Module: + return self._model diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 1c015b464b..f5b11da48e 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -203,6 +203,71 @@ def __hash__(self): ) +# TODO(agagik): change splash_attention_mask._ComputableMask to be non protected +class BlockDiffusionMask(splash_attention_mask._ComputableMask): # pylint: disable=protected-access + """Lazy block-diffusion mask. + + Attention is bidirectional within each block of `bd_size` tokens and causal + across blocks: query position i attends key position j iff + ``(i // bd_size) >= (j // bd_size)``. This converts a causal LM into a + block-diffusion LM (Fast_dLLM / Diffusion-GR2). It is weight-preserving: only + the attention mask changes, not the model weights. + + This mask class inherits from splash_attention_mask._ComputableMask and is + designed to be used with Splash Attention. It computes the mask logic + on-the-fly, avoiding materializing the full (seq_len, seq_len) boolean array. + """ + + #: The size of each bidirectional block. + bd_size: int + + def __init__( + self, + shape: tuple[int, int], + bd_size: int, + shard_count: int = 1, + ): + if bd_size <= 0: + raise ValueError("bd_size must be positive") + self.bd_size = bd_size + + # Define the mask function for block-diffusion attention. + def block_diffusion_mask_function(q_ids, kv_ids): + """Computes the mask logic for the given slice indices.""" + if q_ids.size == 0 or kv_ids.size == 0: + return np.empty((q_ids.shape[0], kv_ids.shape[1]), dtype=np.bool_) + + # Bidirectional within a block, causal across blocks: attend iff the query + # block index is at or after the key block index. + return (q_ids // self.bd_size) >= (kv_ids // self.bd_size) + + # Initialize the parent ComputableMask with this function + super().__init__( + shape=shape, + mask_function=block_diffusion_mask_function, + shard_count=shard_count, + ) + + # Implement equality and hashing based on relevant attributes + def __eq__(self, other: object): + if not isinstance(other, type(self)): + return NotImplemented + # Compare shape, bd_size, and the underlying q_sequence array + return ( + self.shape == other.shape and self.bd_size == other.bd_size and np.array_equal(self.q_sequence, other.q_sequence) + ) + + def __hash__(self): + return hash( + ( + type(self), + self.shape, + self.bd_size, + self.q_sequence.tobytes() if self.q_sequence is not None else None, + ) + ) + + def _generate_chunk_attention_mask(mask_shape: tuple[int, int], chunk_size: int, q_offset: int = 0) -> jax.Array: """Generates an explicit boolean mask for chunked causal attention. @@ -322,6 +387,7 @@ def attention_op_as_linen( attn_logits_soft_cap: float | None = None, sliding_window_size: int | None = None, chunk_attn_window_size: int | None = None, + bd_size: int | None = None, use_ragged_attention: bool = False, ragged_block_size: int = 256, ): @@ -360,6 +426,7 @@ def attention_op_as_linen( attn_logits_soft_cap=attn_logits_soft_cap, sliding_window_size=sliding_window_size, chunk_attn_window_size=chunk_attn_window_size, + bd_size=bd_size, use_ragged_attention=use_ragged_attention, ragged_block_size=ragged_block_size, metadata_fn=variable_to_logically_partitioned, @@ -419,6 +486,7 @@ def __init__( attn_logits_soft_cap: float | None = None, sliding_window_size: int | None = None, chunk_attn_window_size: int | None = None, + bd_size: int | None = None, use_ragged_attention: bool = False, ragged_block_size: int = 256, rngs: nnx.Rngs | None = None, @@ -516,6 +584,7 @@ def __init__( self.attn_logits_soft_cap = attn_logits_soft_cap self.sliding_window_size = sliding_window_size self.chunk_attn_window_size = chunk_attn_window_size + self.bd_size = bd_size self.use_ragged_attention = use_ragged_attention self.ragged_block_size = ragged_block_size self.rngs = rngs @@ -729,6 +798,7 @@ def generate_attention_mask( if model_mode != MODEL_MODE_AUTOREGRESSIVE and self.attention_type not in ( AttentionType.FULL, AttentionType.COMPRESSED, + AttentionType.BLOCK_DIFFUSION, ): if use_segment_positions: causal_mask = (position_col_ids <= position_row_ids)[:, None, None, :, :] @@ -808,6 +878,24 @@ def generate_attention_mask( ) output_mask = chunk_mask * output_mask + elif self.attention_type == AttentionType.BLOCK_DIFFUSION: + if self.bd_size is None: + raise ValueError("bd_size must be set for block_diffusion attention type") + # Block-diffusion: bidirectional within a block of `bd_size` tokens and + # causal across blocks. Query i attends key j iff (i // bd) >= (j // bd), + # i.e. (col_ids // bd) <= (row_ids // bd). The standard causal_mask was + # skipped above (guarded like FULL/COMPRESSED), so it is intentionally NOT + # ANDed in here; the block mask itself provides cross-block causality. + if use_segment_positions: + block_mask = ((position_col_ids // self.bd_size) <= (position_row_ids // self.bd_size))[:, None, None, :, :] + else: + mask_shape = (q_seq_len, kv_seq_len) + row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0) + next_pos + col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1) + block_mask = ((col_ids // self.bd_size) <= (row_ids // self.bd_size))[None, None, None, :, :] + # Combine only with the segment/packing mask (if present), never with causal. + output_mask = block_mask if output_mask is None else (block_mask & output_mask) + if bidirectional_mask is not None: image_mask = _make_bidirectional_block_mask(bidirectional_mask) output_mask = output_mask | image_mask[:, None, None, ...] @@ -1301,11 +1389,17 @@ def create_sa_config(config, query, key, attn_logits_soft_cap): mask_module = tokamax_splash_mask if self.config.use_tokamax_splash else splash_attention_mask if self.attention_type == AttentionType.FULL: mask = mask_module.FullMask(mask_shape) + elif self.attention_type == AttentionType.BLOCK_DIFFUSION: + if self.bd_size is None: + raise ValueError("bd_size must be set for block_diffusion attention type") + # Base mask is the block-diffusion mask itself (bidirectional within block, + # causal across blocks). Do NOT intersect with CausalMask below. + mask = BlockDiffusionMask(shape=mask_shape, bd_size=self.bd_size) else: mask = mask_module.CausalMask(shape=mask_shape) use_load_balanced_cp = cp_size > 1 and load_balanced_context_parallel - if use_load_balanced_cp and self.attention_type != AttentionType.FULL: + if use_load_balanced_cp and self.attention_type not in (AttentionType.FULL, AttentionType.BLOCK_DIFFUSION): mask = LoadBalancedCausalMask(shape=mask_shape, cp_size=cp_size) # Apply local masking if local sliding attention is enabled. diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index b08a059938..d219de7507 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -462,6 +462,7 @@ def __init__( attn_logits_soft_cap=self.attn_logits_soft_cap, sliding_window_size=self.sliding_window_size, chunk_attn_window_size=self.config.chunk_attn_window_size, + bd_size=self.config.bd_size, use_ragged_attention=self.use_ragged_attention, ragged_block_size=self.ragged_block_size, rngs=self.rngs, diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py index 3b5f08f0c1..ef8a1c24da 100644 --- a/src/maxtext/models/qwen3.py +++ b/src/maxtext/models/qwen3.py @@ -1407,6 +1407,7 @@ def __init__( max_target_length=config.max_target_length, max_prefill_predict_length=config.max_prefill_predict_length, attention_kernel=config.attention, + attention_type=AttentionType.BLOCK_DIFFUSION if config.enable_block_diffusion else AttentionType.GLOBAL, inputs_q_shape=dummy_inputs_shape, inputs_kv_shape=dummy_inputs_shape, mesh=mesh, diff --git a/src/maxtext/trainers/post_train/distillation/diffusion_generate.py b/src/maxtext/trainers/post_train/distillation/diffusion_generate.py new file mode 100644 index 0000000000..9681627f1a --- /dev/null +++ b/src/maxtext/trainers/post_train/distillation/diffusion_generate.py @@ -0,0 +1,364 @@ +# Copyright 2023–2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Student block-diffusion rollout for on-policy distillation (OPD). + +This is the MaxText training-side analog of tpu_inference's standalone +``runner/diffusion_decode.py`` inference loop. The student is a block-diffusion +model (``enable_block_diffusion`` on): attention is bidirectional within a block +of ``bd_size`` tokens and block-causal across blocks. Given a prompt prefix, the +student generates its own completion by block-diffusion *denoising*, so the loss +can be computed on-policy against a frozen teacher (see +``train_distill.MaxTextDistillationTrainer._run_on_policy_rollout``). + +Design +------ +The module is intentionally free of MaxText / Tunix imports (only ``jax``) so the +whole seed -> forward -> shift -> commit -> advance loop is exercisable on CPU +with a stub ``forward_fn`` (dependency injection), exactly like the reference +inference loop. The real student is injected by the trainer as:: + + forward_fn(tokens[B, T] int32, positions[B, T] int32) -> logits[B, T, V] + +Unlike the inference reference (which forwards a single ``block_size`` canvas +against a paged KV cache), the MaxText trainer forward is *cacheless* and returns +per-position logits for the whole ``[B, T]`` sequence. We therefore operate on +the full sequence and rely on the block-causal attention mask: when denoising +block ``k`` the model's logits for block ``k`` depend only on blocks ``0..k``, so +the still-masked trailing blocks never corrupt an earlier block and we can fill +strictly left-to-right, one block at a time. + +Conventions (matching the reference decode loop) +------------------------------------------------ +* **Seed / canvas**: the completion positions (``gen_mask``) are seeded with + ``mask_id``; the prompt positions keep their real tokens and are never masked. +* **Shifted logits**: the token at position ``i`` is predicted from hidden + ``i - 1`` (clamped ``>= 0``), i.e. ``shifted[:, i] = logits[:, max(i - 1, 0)]`` + — the AR next-token convention. For the first completion position this reads + the last prompt hidden, so no separate "first token" seed is needed. +* **Threshold-commit + forced argmax**: each denoise iteration commits the + masked positions whose top-1 probability exceeds ``threshold`` and always + force-commits the single highest-confidence still-masked position per row, so + a block of ``bd_size`` positions fully commits in ``<= bd_size`` iterations. + +Loop design +----------- +The **inner** per-block denoise iterations run as a ``jax.lax.while_loop`` and +the **outer** block advance runs as a ``jax.lax.fori_loop`` (sequence length is +fixed in training, so the block count ``T // bd_size`` is static) — the whole +rollout is one jittable, ``stop_gradient``-able program. Generation is OFF the +student gradient path: the trainer runs this outside ``value_and_grad`` and wraps +the outputs in ``jax.lax.stop_gradient`` (the committed tokens are argmax indices, +which are non-differentiable regardless). +""" + +from typing import Callable, Tuple + +import jax +import jax.numpy as jnp + +# forward_fn(tokens[B, T] int32, positions[B, T] int32) -> logits[B, T, V] +ForwardFn = Callable[[jax.Array, jax.Array], jax.Array] + + +def _shifted_logit_indices(seq_len: int) -> jax.Array: + """Indices implementing the shifted-logit convention. + + ``shifted[:, i] = logits[:, max(i - 1, 0)]`` so the token at position ``i`` is + predicted from hidden ``i - 1`` (clamped ``>= 0`` for position 0). + """ + return jnp.maximum(jnp.arange(seq_len, dtype=jnp.int32) - 1, 0) + + +def _selective_log_softmax(logits: jax.Array, tokens: jax.Array) -> jax.Array: + """Log-prob of ``tokens`` under ``logits`` (pure-jax; no tunix dependency). + + ``logits`` is ``(..., L, V)`` and ``tokens`` is ``(..., L)``; returns ``(..., L)`` + where entry ``p`` is ``log softmax(logits[..., p, :])[tokens[..., p]]``. Mirrors + ``tunix.rl.common.selective_log_softmax`` so the diffusion logprob and the AR + logprob agree numerically on the extraction step. + """ + logits = logits.astype(jnp.float32) + target = jnp.take_along_axis(logits, tokens[..., None], axis=-1).squeeze(-1) + normalizer = jax.nn.logsumexp(logits, axis=-1) + return target - normalizer + + +def _positions_from_mask(mask: jax.Array) -> jax.Array: + """Absolute positions from a validity mask (RoPE), pure-jax. + + Reimplements ``tunix.sft.utils.build_positions_from_mask`` + (``cumsum(mask) - (cumsum(mask) >= 1)``) so this module stays jax-only: pad + positions and the first real token share position ``0`` and real tokens then + count ``0, 1, 2, ...``. Block-diffusion blocks are ``position // bd_size``, so + using the SAME position convention as the AR path keeps the two logprob paths + aligned. + """ + positions = jnp.cumsum(mask.astype(jnp.int32), axis=-1) + return positions - (positions >= 1).astype(jnp.int32) + + +def diffusion_commit( + logits: jax.Array, + mask: jax.Array, + threshold: float, + temperature: float = 0.0, +) -> Tuple[jax.Array, jax.Array]: + """Per-position, threshold-based commit step for block-diffusion denoising. + + MaxText-side reimplementation of ``tpu_inference``'s ``diffusion_commit`` (that + repo is a separate, read-only checkout, so the pure-jax logic is duplicated + here rather than imported). Given per-position vocab logits and a boolean mask + marking positions that are still masked (not yet committed), greedily commit + the positions whose top-1 probability exceeds ``threshold``. To guarantee + forward progress, the single highest-confidence still-masked position in each + row is always committed even if it is below ``threshold`` (unless the row has + no masked positions left). + + Args: + logits: ``(..., L, V)`` float logits. Any number of leading batch/row dims is + supported (e.g. ``(B, T, V)`` or ``(L, V)``). + mask: ``(..., L)`` boolean array; ``True`` marks positions still masked and + therefore eligible to be committed this step. + threshold: scalar confidence threshold in ``[0, 1]``. A masked position + commits when its top-1 softmax probability is strictly greater than this. + temperature: optional softmax temperature; when ``> 0`` the logits are + divided by it before the softmax, when ``0`` (default) the raw logits are + used. + + Returns: + ``(committed_token_ids, new_mask)`` where ``committed_token_ids`` is the + ``(..., L)`` int32 top-1 token id for every position (values at still-masked + positions should be ignored by the caller), and ``new_mask`` equals ``mask`` + with the newly-committed positions cleared (a strictly shrinking mask). + """ + if temperature > 0.0: + logits = logits / temperature + probs = jax.nn.softmax(logits.astype(jnp.float32), axis=-1) + + top_prob = jnp.max(probs, axis=-1) + top_tok = jnp.argmax(probs, axis=-1).astype(jnp.int32) + + mask_bool = mask.astype(bool) + + # Threshold commit: only sufficiently-confident masked positions commit. + commit = (top_prob > threshold) & mask_bool + + # Progress guarantee: force the single highest-confidence still-masked position + # in each row to commit. Non-masked positions are excluded from the argmax so + # an already-committed position can never be re-selected, and rows with no + # masked positions left force nothing. + neg_inf = jnp.array(-jnp.inf, dtype=top_prob.dtype) + masked_conf = jnp.where(mask_bool, top_prob, neg_inf) + forced_idx = jnp.argmax(masked_conf, axis=-1) + row_has_mask = jnp.any(mask_bool, axis=-1) + forced_onehot = jax.nn.one_hot(forced_idx, mask_bool.shape[-1], dtype=bool) + forced_onehot = forced_onehot & jnp.expand_dims(row_has_mask, axis=-1) + commit = commit | forced_onehot + + new_mask = mask_bool & (~commit) + return top_tok, new_mask + + +def diffusion_generate( + forward_fn: ForwardFn, + tokens: jax.Array, + positions: jax.Array, + gen_mask: jax.Array, + *, + mask_id: int, + bd_size: int, + threshold: float = 0.9, + temperature: float = 0.0, + max_denoise_steps: int = 0, +) -> Tuple[jax.Array, jax.Array]: + """Generate a completion for a batch by block-diffusion denoising. + + Fills the ``gen_mask`` positions of ``tokens`` block-by-block (blocks aligned + to ``bd_size``, matching the block-diffusion attention's ``position // bd_size``) + using the injected student ``forward_fn``. Prompt positions (``gen_mask`` False) + are held fixed. + + Args: + forward_fn: injected cacheless student forward. See module docstring for the + signature ``(tokens, positions) -> logits[B, T, V]``. + tokens: ``[B, T]`` int32 batch; prompt region holds real tokens, the + completion region is (re)seeded with ``mask_id`` internally. + positions: ``[B, T]`` int32 absolute positions (for RoPE), passed through to + ``forward_fn``. + gen_mask: ``[B, T]`` boolean; ``True`` where the student should generate. + mask_id: token id used to mark not-yet-committed canvas positions. + bd_size: block size; ``T`` must be divisible by it. + threshold: confidence threshold for committing a position. + temperature: softmax temperature for the commit step (``0.0`` = argmax). + max_denoise_steps: per-block denoise iteration cap; ``<= 0`` means use + ``bd_size`` (which fully commits, since the forced-argmax progress guarantee + commits ``>= 1`` masked position per iteration). + + Returns: + ``(gen_tokens, committed_mask)`` where ``gen_tokens`` is the ``[B, T]`` int32 + sequence (prompt ⊕ generated completion, no ``mask_id`` left) and + ``committed_mask`` is the ``[B, T]`` boolean mask of the positions the student + generated (equals ``gen_mask`` once every block is filled). + """ + tokens = jnp.asarray(tokens, dtype=jnp.int32) + positions = jnp.asarray(positions, dtype=jnp.int32) + gen_mask = jnp.asarray(gen_mask, dtype=bool) + + batch_size, seq_len = tokens.shape + if bd_size <= 0: + raise ValueError(f"bd_size must be > 0, got {bd_size}") + if seq_len % bd_size != 0: + raise ValueError(f"sequence length ({seq_len}) must be divisible by bd_size ({bd_size})") + num_blocks = seq_len // bd_size + eff_steps = bd_size if max_denoise_steps <= 0 else int(max_denoise_steps) + + shift_idx = _shifted_logit_indices(seq_len) # [T] + pos_ids = jnp.arange(seq_len, dtype=jnp.int32) # [T] + + # Seed canvas: mask_id at generation positions, real tokens (prompt) elsewhere. + canvas0 = jnp.where(gen_mask, jnp.array(mask_id, dtype=jnp.int32), tokens) # [B, T] + committed0 = jnp.zeros((batch_size, seq_len), dtype=bool) + + def block_body(blk, carry): + canvas, committed = carry + start = blk * bd_size + in_block = (pos_ids >= start) & (pos_ids < start + bd_size) # [T] + block_gen = gen_mask & in_block[None, :] # [B, T] positions to denoise this block + has_gen = jnp.any(block_gen) + + def cond_fn(inner): + _, mask, step, _ = inner + under_cap = step < eff_steps + # Always run step 0 so a degenerate all-committed block still terminates; + # otherwise stop once every eligible position in the block is committed. + more_work = jnp.logical_or(step == 0, jnp.any(mask)) + return jnp.logical_and(under_cap, more_work) + + def body_fn(inner): + cv, mask, step, _ = inner + logits = forward_fn(cv, positions).astype(jnp.float32) # [B, T, V] + shifted = logits[:, shift_idx, :] # [B, T, V] — token i predicted from hidden i-1 + committed_tok, new_mask = diffusion_commit(shifted, mask, threshold, temperature) + newly = jnp.logical_and(mask, jnp.logical_not(new_mask)) + cv = jnp.where(newly, committed_tok, cv) + return (cv, new_mask, step + 1, committed_tok) + + def do_denoise(cv): + init = (cv, block_gen, jnp.array(0, dtype=jnp.int32), cv) + out_canvas, mask_final, _, last_tok = jax.lax.while_loop(cond_fn, body_fn, init) + # If the step cap was hit before full commitment, force-fill any still-masked + # positions with their argmax so mask_id never leaks into the output. + return jnp.where(mask_final, last_tok, out_canvas) + + # Skip the (wasted) forward for prompt-only blocks with nothing to generate. + canvas = jax.lax.cond(has_gen, do_denoise, lambda cv: cv, canvas) + committed = committed | block_gen + return (canvas, committed) + + canvas, committed = jax.lax.fori_loop(0, num_blocks, block_body, (canvas0, committed0)) + return canvas, committed + + +def diffusion_per_token_logps( + model_apply: ForwardFn, + prompt_tokens: jax.Array, + completion_tokens: jax.Array, + completion_mask: jax.Array | None = None, + *, + bd_size: int, + pad_id: int = 0, + temperature: float = 1.0, + stop_gradient: bool = False, +) -> jax.Array: + """Block-diffusion (TraceRL-style) per-token log-probability of a completion. + + THE single source of truth for the GRPO/DAPO importance ratio when the policy + is a block-diffusion model. It is invoked at BOTH sites of the ratio — the + rollout's old-policy logps (``MaxTextDiffusionRollout.get_per_token_logps``) + and the loss's new-policy logps (the tunix ``per_token_logps_fn`` hook) — so the + two are guaranteed identical and the ratio is unbiased. Do NOT reimplement the + math at either call site; both must funnel through this function. + + Unlike an autoregressive next-token logprob, the score comes from a single + block-diffusion forward over the whole committed sequence: each position attends + bidirectionally within its block of ``bd_size`` tokens and block-causally across + blocks (the model applies this mask internally when ``enable_block_diffusion`` is + on; ``model_apply`` just forwards the full sequence). The completed sequence is + then scored with the same shifted-logit convention the rollout generated it under + (``_shifted_logit_indices``: ``shifted[:, p] = logits[:, p - 1]``, so the + post-shift logits at position ``p`` predict token ``p`` — matching + ``diffusion_generate``'s commit rule, which argmaxes ``logits[p - 1]`` onto token + ``p``). Scoring under the generation rule is what makes the old-policy logp equal + the probability the token was drawn with. + + Args: + model_apply: cacheless block-diffusion forward ``(tokens[B, T], positions[B, + T]) -> logits[B, T, V]`` (same signature/injection as ``diffusion_generate``, + so it is exercisable on CPU with a stub). + prompt_tokens: ``[B, P]`` int prompt ids (may be left-padded with ``pad_id``). + completion_tokens: ``[B, C]`` int completion ids (may be right-padded). + completion_mask: optional ``[B, C]`` mask; when provided the returned logps are + zeroed outside it and it also defines completion validity for positions. + bd_size: block size; ``P + C`` must be divisible by it (block-diffusion + contract, matching ``diffusion_generate`` and the config validator). + pad_id: pad id used to derive the prompt validity mask for positions. + temperature: divides the logits before the softmax when not ``0.0``/``1.0`` + (matches ``tunix.rl.common.compute_per_token_logps``). + stop_gradient: when True the result is wrapped in ``jax.lax.stop_gradient`` + (used for the rollout's old-policy logps; the loss passes ``False`` so the + policy gradient flows). + + Returns: + ``[B, C]`` per-token log-probabilities for the completion positions. + """ + prompt_tokens = jnp.asarray(prompt_tokens, dtype=jnp.int32) + completion_tokens = jnp.asarray(completion_tokens, dtype=jnp.int32) + if bd_size <= 0: + raise ValueError(f"bd_size must be > 0, got {bd_size}") + + prompt_len = prompt_tokens.shape[1] + full_tokens = jnp.concatenate([prompt_tokens, completion_tokens], axis=1) # [B, T] + seq_len = full_tokens.shape[1] + if seq_len % bd_size != 0: + raise ValueError(f"prompt+completion length ({seq_len}) must be divisible by bd_size ({bd_size})") + + # Validity mask -> RoPE positions (block index = position // bd_size). Prompt + # validity comes from pad_id; completion validity from completion_mask (or + # non-pad when absent), matching the AR path so both logprob paths align. + prompt_mask = prompt_tokens != pad_id + if completion_mask is None: + completion_valid = completion_tokens != pad_id + else: + completion_valid = completion_mask.astype(bool) + full_mask = jnp.concatenate([prompt_mask, completion_valid], axis=1) + positions = _positions_from_mask(full_mask) # [B, T] + + # Single block-diffusion forward over the committed sequence. + logits = model_apply(full_tokens, positions).astype(jnp.float32) # [B, T, V] + + # Shifted-logit convention (reuse M3): post-shift logits at p predict token p. + shift_idx = _shifted_logit_indices(seq_len) # [T] + shifted = logits[:, shift_idx, :] # [B, T, V] + if temperature != 0.0 and temperature != 1.0: + shifted = shifted / temperature + + logps_full = _selective_log_softmax(shifted, full_tokens) # [B, T] + completion_logps = logps_full[:, prompt_len:] # [B, C] + + if completion_mask is not None: + completion_logps = completion_logps * completion_mask.astype(completion_logps.dtype) + if stop_gradient: + completion_logps = jax.lax.stop_gradient(completion_logps) + return completion_logps diff --git a/src/maxtext/trainers/post_train/distillation/train_distill.py b/src/maxtext/trainers/post_train/distillation/train_distill.py index 15197d1af4..7c8cd2f2a0 100644 --- a/src/maxtext/trainers/post_train/distillation/train_distill.py +++ b/src/maxtext/trainers/post_train/distillation/train_distill.py @@ -55,7 +55,7 @@ from maxtext.input_pipeline import input_pipeline_interface from maxtext.layers.learn_to_init_layer import apply_lti_model_update from maxtext.optimizers import optimizers -from maxtext.trainers.post_train.distillation import distillation_utils, lti_utils +from maxtext.trainers.post_train.distillation import diffusion_generate, distillation_utils, lti_utils from maxtext.utils import max_logging from maxtext.utils import maxtext_utils from maxtext.utils import model_creation_utils @@ -234,6 +234,28 @@ def __init__( self.strategy = strategy self.checkpoint_manager: distillation_utils.MaxTextCheckpointManager = None + # On-policy distillation (OPD) for block-diffusion. When enabled, the student + # generates its own completion by block-diffusion denoising before the loss is + # computed (see `_run_on_policy_rollout`); when disabled the trainer is the + # standard (off-policy) distiller and `_train_step` is unchanged. + self.opd_on_policy = getattr(student_config, "opd_on_policy", False) + self.opd_bd_size = getattr(student_config, "bd_size", 32) + self.opd_mask_id = getattr(student_config, "mask_id", 151669) + self.opd_threshold = getattr(student_config, "opd_threshold", 0.9) + self.opd_temperature = getattr(student_config, "opd_temperature", 0.0) + self.opd_max_denoise_steps = getattr(student_config, "opd_max_denoise_steps", 0) + if self.opd_on_policy: + max_logging.log( + "On-policy distillation (OPD) ENABLED: block-diffusion student rollout " + f"(bd_size={self.opd_bd_size}, mask_id={self.opd_mask_id}, threshold={self.opd_threshold}, " + f"temperature={self.opd_temperature}, max_denoise_steps={self.opd_max_denoise_steps})." + ) + if not getattr(student_config, "enable_block_diffusion", False): + max_logging.log( + "WARNING: opd_on_policy=True but enable_block_diffusion=False; the student rollout assumes " + "block-diffusion attention. Set enable_block_diffusion=True for the student." + ) + # Per-step per-device TFLOPs (constants for the run): student fwd+bwd + teacher fwd-only. ( self._tflops_combined, @@ -290,6 +312,15 @@ def _train_step(self, model, optimizer, inputs): teacher = model.teacher_model current_step = model.training_step[...] + # On-policy distillation (OPD): the student generates its own completion by + # block-diffusion denoising and the batch is rewritten so both teacher and + # student are scored on that generated sequence, with the forward-KL applied + # only at the committed (generated) positions. Run BEFORE teacher_output and + # OUTSIDE value_and_grad, exactly like the frozen teacher — generation is off + # the student gradient path. + if self.opd_on_policy: + batch = self._run_on_policy_rollout(batch, student) + # Run teacher inference outside of value_and_grad. # The teacher is frozen (stop_gradient), so its output is a constant # from the perspective of the student gradient computation. @@ -345,6 +376,79 @@ def loss_wrapper_pure(diff_params, rest): return loss, aux, optax.global_norm(grads) return loss, aux + def _run_on_policy_rollout(self, batch: dict, student: nnx.Module) -> dict: + """Rewrites `batch` with the student's own on-policy block-diffusion completion. + + The student is a block-diffusion model (`enable_block_diffusion` on). Given the + batch prompts, it denoises its own completion via + `diffusion_generate.diffusion_generate` (seed canvas -> student forward -> + shifted logits -> threshold-commit + forced argmax -> advance block by block). + The batch is then rewritten so that: + + * `input_tokens` / `targets` = prompt ⊕ student_completion (aligned), so both + the frozen teacher and the student are scored on the generated sequence; + * `targets_segmentation` = the committed-position mask, so the reused + forward-KL loss (`compute_loss`) applies ONLY at the committed positions. + + Generation is OFF the student gradient path: it is run here (outside + `value_and_grad`, before the teacher forward) and the outputs are wrapped in + `jax.lax.stop_gradient`. The committed tokens are argmax indices, which are + non-differentiable regardless. + + Args: + batch: the input dict produced by `gen_model_input_fn`. + student: the (block-diffusion) student module. + + Returns: + A new batch dict with `input_tokens`, `targets`, and `targets_segmentation` + replaced by the on-policy rollout. + """ + input_tokens = batch["input_tokens"] + positions = batch["positions"] + decoder_segment_ids = batch.get("decoder_segment_ids") + + # Completion span = where the student should generate. Prompt (and padding) + # carry targets_segmentation == 0 and are held fixed. Fall back to non-pad + # tokens when no segmentation is available. + seg = batch.get("targets_segmentation") + if seg is not None: + gen_mask = seg != 0 + else: + gen_mask = input_tokens != self.strategy.pad_id + + def gen_forward_fn(tokens, pos): + # Direct cacheless student forward; block-diffusion attention is active via + # enable_block_diffusion. No dropout / feature / MoE collection during rollout. + return student( + decoder_input_tokens=tokens, + decoder_positions=pos, + decoder_segment_ids=decoder_segment_ids, + enable_dropout=False, + ) + + gen_tokens, committed_mask = diffusion_generate.diffusion_generate( + gen_forward_fn, + input_tokens, + positions, + gen_mask, + mask_id=self.opd_mask_id, + bd_size=self.opd_bd_size, + threshold=self.opd_threshold, + temperature=self.opd_temperature, + max_denoise_steps=self.opd_max_denoise_steps, + ) + # Generation is a constant w.r.t. the student gradient. + gen_tokens = jax.lax.stop_gradient(gen_tokens) + committed_mask = jax.lax.stop_gradient(committed_mask) + + new_batch = dict(batch) + new_batch["input_tokens"] = gen_tokens + # Aligned targets = the student-generated tokens; the committed-position mask + # (written into targets_segmentation) is what gates the forward-KL loss. + new_batch["targets"] = gen_tokens + new_batch["targets_segmentation"] = committed_mask.astype(input_tokens.dtype) + return new_batch + def _eval_step(self, model, inputs): """Evaluation only needs the student.""" inputs = self.gen_model_input_fn(inputs) diff --git a/src/maxtext/trainers/post_train/rl/train_rl.py b/src/maxtext/trainers/post_train/rl/train_rl.py index 19f714f7ae..16cb6b12a9 100644 --- a/src/maxtext/trainers/post_train/rl/train_rl.py +++ b/src/maxtext/trainers/post_train/rl/train_rl.py @@ -117,6 +117,10 @@ def _no_bf16_to_f32_cast(val, tgt_dtype, src_key): from maxtext.configs import pyconfig from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR from maxtext.integration.vllm.maxtext_vllm_rollout import MaxTextVllmRollout +from maxtext.integration.vllm.maxtext_diffusion_rollout import ( + MaxTextDiffusionRollout, + make_diffusion_per_token_logps_fn, +) from maxtext.trainers.post_train.rl.evaluate_rl import evaluate from maxtext.trainers.post_train.rl import utils_rl from maxtext.input_pipeline.instruction_data_processing import load_data_template_from_file @@ -450,11 +454,25 @@ def create_rl_components( argv_list = ["", str(vllm_config_path), "log_config=False"] vllm_config = pyconfig.initialize(argv_list) - rl_rollout_engine = ( - functools.partial(MaxTextVllmRollout, maxtext_config=trainer_config) - if trainer_config.use_standalone_converter - else "vllm" - ) + # Block-diffusion RL (Fast_dLLM / Diffusion-GR2, TraceRL-style): the policy + # generates by block-diffusion denoising and the importance ratio uses the shared + # block-diffusion per-token logprob. Gated behind `rl.rl_diffusion_rollout` so AR + # RL (vLLM or standalone converter) is entirely unchanged when it is off. The + # matching new-policy logprob is injected into the GRPO learner below via + # `per_token_logps_fn` so both sides of the ratio use the SAME function. + use_diffusion_rollout = getattr(trainer_config.rl, "rl_diffusion_rollout", False) + if use_diffusion_rollout: + if not trainer_config.enable_block_diffusion: + raise ValueError( + "rl.rl_diffusion_rollout=True requires enable_block_diffusion=True: the " + "diffusion rollout denoises a block-diffusion canvas and scores it with " + "block-diffusion attention." + ) + rl_rollout_engine = functools.partial(MaxTextDiffusionRollout, maxtext_config=trainer_config) + elif trainer_config.use_standalone_converter: + rl_rollout_engine = functools.partial(MaxTextVllmRollout, maxtext_config=trainer_config) + else: + rl_rollout_engine = "vllm" cluster_config = rl_cluster_lib.ClusterConfig( role_to_mesh={ @@ -608,10 +626,20 @@ def _reward_fn(**kwargs): loss_algo=trainer_config.rl.loss_algo, loss_agg_mode=trainer_config.rl.loss_agg_mode, ) + # When the diffusion rollout is active, override the new-policy per-token + # logprob with the SAME shared block-diffusion fn the rollout uses for the + # old-policy logps (the P0 invariant); otherwise the learner keeps the default + # AR `common.compute_per_token_logps` and AR RL is unchanged. + diffusion_per_token_logps_fn = ( + make_diffusion_per_token_logps_fn(trainer_config.bd_size, trainer_config.mask_id) + if use_diffusion_rollout + else None + ) rl_trainer = GrpoLearner( rl_cluster=rl_cluster, reward_fns=reward_fns, algo_config=grpo_config, + per_token_logps_fn=diffusion_per_token_logps_fn, ) return rl_cluster, rl_trainer, optimizer, reward_fns diff --git a/tests/post_training/unit/opd_generate_test.py b/tests/post_training/unit/opd_generate_test.py new file mode 100644 index 0000000000..ec0b5718ed --- /dev/null +++ b/tests/post_training/unit/opd_generate_test.py @@ -0,0 +1,394 @@ +# Copyright 2023–2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU unit tests for on-policy distillation (OPD) block-diffusion rollout. + +Covers the *pure* logic of the OPD student rollout end-to-end on CPU: + + (a) the student block-diffusion generate loop with a STUB forward (canned, + position-addressed logits): a block fills within its step cap, the + forced-argmax progress guarantee, the shifted-logit indexing (token i from + hidden i-1), multi-block advancement, prompt positions held fixed, and the + step-cap force-fill (no ``mask_id`` leak); + (b) that the committed-position mask makes the reused forward-KL loss nonzero + ONLY at committed positions (real ``create_labels`` + ``compute_loss`` with + crafted student/teacher logits); + (c) that generation is stop-gradient — no gradient flows to the (differentiable) + student forward parameters through the rollout (checked via ``jax.grad``), + plus a source check that ``_train_step`` wires the rollout under + ``jax.lax.stop_gradient`` and behind the ``opd_on_policy`` gate. + +The pure rollout module (``diffusion_generate.py``) is loaded standalone by file +path (it imports only ``jax``), so (a)/(c) run with no MaxText/Tunix deps. The +loss gating test (b) imports the real ``distillation_utils``; when ``tunix`` is +not installed (bare CPU venv) a minimal stub is injected so the real strategy +still loads and runs. The full two-model on-policy ``_train_step`` needs real +model weights and a TPU and is therefore covered by source/AST checks here. +""" + +import ast +import importlib.util +import pathlib +import sys +import types +from typing import Any + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +pytestmark = [pytest.mark.cpu_only, pytest.mark.post_training] + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +_DISTILL_DIR = _REPO_ROOT / "src" / "maxtext" / "trainers" / "post_train" / "distillation" +_DIFFUSION_GENERATE = _DISTILL_DIR / "diffusion_generate.py" +_TRAIN_DISTILL = _DISTILL_DIR / "train_distill.py" + + +def _load_diffusion_generate_standalone(): + """Load ``diffusion_generate.py`` by file path (it imports only jax).""" + spec = importlib.util.spec_from_file_location("opd_diffusion_generate_under_test", _DIFFUSION_GENERATE) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +_DG = _load_diffusion_generate_standalone() +diffusion_generate = _DG.diffusion_generate +diffusion_commit = _DG.diffusion_commit + + +def _import_distillation_utils(): + """Import the real ``distillation_utils``; stub ``tunix`` if it is absent. + + In CI ``tunix`` is installed and imported directly. On a bare ``jax[cpu]`` venv + we inject the two tiny tunix symbols ``distillation_utils`` needs at import time + (``peft_trainer.TrainingInput`` / ``PeftTrainer`` and + ``checkpoint_manager.CheckpointManager``) so the real strategy class still loads. + """ + try: + import tunix # noqa: F401 pylint: disable=unused-import + except ImportError: + import flax + + for pkg in ("tunix", "tunix.sft"): + if pkg not in sys.modules: + m = types.ModuleType(pkg) + m.__path__ = [] # mark as package + sys.modules[pkg] = m + + pt = types.ModuleType("tunix.sft.peft_trainer") + + @flax.struct.dataclass(frozen=True) + class _TrainingInput: # minimal base for MaxTextTrainingInput + input_tokens: Any + input_mask: Any + + pt.TrainingInput = _TrainingInput + pt.PeftTrainer = type("PeftTrainer", (), {}) + sys.modules["tunix.sft.peft_trainer"] = pt + + cm = types.ModuleType("tunix.sft.checkpoint_manager") + cm.CheckpointManager = type("CheckpointManager", (), {}) + sys.modules["tunix.sft.checkpoint_manager"] = cm + + src = str(_REPO_ROOT / "src") + if src not in sys.path: + sys.path.insert(0, src) + from maxtext.trainers.post_train.distillation import distillation_utils # pylint: disable=import-outside-toplevel + + return distillation_utils + + +# ---- Test constants --------------------------------------------------------- +VOCAB = 16 +MASK_ID = 15 # out of the normal token range so leaks are detectable + + +def _target_array(n: int) -> np.ndarray: + """Deterministic absolute-position -> token map in [0, VOCAB-3].""" + return np.array([p % (VOCAB - 2) for p in range(n)], dtype=np.int32) + + +def _make_stub_forward(global_target: np.ndarray, logit_scale: float): + """Pure-JAX stub forward ``(tokens[B,T], positions[B,T]) -> logits[B,T,V]``. + + ``logits[b, j]`` peaks on ``global_target[positions[b, j]]`` with confidence + set by ``logit_scale`` (large -> commits via threshold in one step; small -> + below threshold so only the forced-argmax position commits). The stub ignores + the canvas so per-position argmax is stable and committed values are exact. + """ + g = jnp.asarray(global_target) + + def forward_fn(tokens, positions): + del tokens + target = g[positions] # [B, T] + return jax.nn.one_hot(target, VOCAB, dtype=jnp.float32) * logit_scale # [B, T, V] + + return forward_fn + + +def _make_batch(batch_size, seq_len, prompt_len): + """Prompt occupies ``[0, prompt_len)``; completion ``[prompt_len, seq_len)``.""" + g = _target_array(256) + tokens = jnp.zeros((batch_size, seq_len), dtype=jnp.int32) + # Distinct, non-mask prompt tokens so "prompt untouched" is meaningful. + prompt = jnp.arange(1, prompt_len + 1, dtype=jnp.int32) + tokens = tokens.at[:, :prompt_len].set(prompt) + positions = jnp.broadcast_to(jnp.arange(seq_len, dtype=jnp.int32), (batch_size, seq_len)) + gen_mask = jnp.broadcast_to(jnp.arange(seq_len) >= prompt_len, (batch_size, seq_len)) + return g, tokens, positions, gen_mask + + +class TestDiffusionGenerateStub: + """(a) Pure student generate loop with a stub forward.""" + + def test_block_fills_within_max_steps_high_confidence(self): + g, tokens, positions, gen_mask = _make_batch(batch_size=2, seq_len=8, prompt_len=4) + fwd = _make_stub_forward(g, logit_scale=30.0) + gen, committed = diffusion_generate( + fwd, tokens, positions, gen_mask, mask_id=MASK_ID, bd_size=4, threshold=0.9, temperature=0.0 + ) + gen = np.asarray(gen) + committed = np.asarray(committed) + # High confidence -> the whole completion commits, no mask_id leaks. + assert not np.any(gen == MASK_ID) + # Every generated position is reported committed; prompt positions are not. + assert np.array_equal(committed, np.asarray(gen_mask)) + + def test_forced_argmax_progress_low_confidence(self): + # Below-threshold confidence: only the forced (highest-confidence masked) + # position commits per step, yet the block still fully fills within bd_size. + g, tokens, positions, gen_mask = _make_batch(batch_size=1, seq_len=8, prompt_len=4) + fwd = _make_stub_forward(g, logit_scale=0.05) + gen, _ = diffusion_generate( + fwd, tokens, positions, gen_mask, mask_id=MASK_ID, bd_size=4, threshold=0.9, temperature=0.0 + ) + assert not np.any(np.asarray(gen) == MASK_ID) + + def test_shifted_logit_indexing_i_from_i_minus_1(self): + # Completion starts mid-block to exercise the shift across a block boundary. + prompt_len = 3 + g, tokens, positions, gen_mask = _make_batch(batch_size=1, seq_len=8, prompt_len=prompt_len) + fwd = _make_stub_forward(g, logit_scale=30.0) + gen, _ = diffusion_generate( + fwd, tokens, positions, gen_mask, mask_id=MASK_ID, bd_size=4, threshold=0.9, temperature=0.0 + ) + gen = np.asarray(gen)[0] + gnp = np.asarray(g) + # Prompt untouched. + assert gen[:prompt_len].tolist() == list(range(1, prompt_len + 1)) + # Shifted convention: committed token at position p == target at (p - 1). + for p in range(prompt_len, 8): + assert int(gen[p]) == int(gnp[p - 1]), f"shift mismatch at {p}" + + def test_multi_block_advance_contiguous_across_boundaries(self): + # prompt block 0 fully prompt; blocks 1 and 2 fully generated. + prompt_len = 4 + seq_len = 12 + bd = 4 + g, tokens, positions, gen_mask = _make_batch(batch_size=1, seq_len=seq_len, prompt_len=prompt_len) + fwd = _make_stub_forward(g, logit_scale=30.0) + gen, committed = diffusion_generate( + fwd, tokens, positions, gen_mask, mask_id=MASK_ID, bd_size=bd, threshold=0.9, temperature=0.0 + ) + gen = np.asarray(gen)[0] + gnp = np.asarray(g) + assert not np.any(gen == MASK_ID) + # Contiguity across every block boundary proves later blocks were seeded from + # the previous block's committed tail (shifted convention, full sequence). + for p in range(prompt_len, seq_len): + assert int(gen[p]) == int(gnp[p - 1]), f"mismatch at {p}" + # Explicit boundary check at the start of block 2. + assert int(gen[2 * bd]) == int(gnp[2 * bd - 1]) + assert np.array_equal(np.asarray(committed)[0], np.asarray(gen_mask)[0]) + + def test_prompt_positions_never_masked(self): + prompt_len = 5 + g, tokens, positions, gen_mask = _make_batch(batch_size=3, seq_len=8, prompt_len=prompt_len) + fwd = _make_stub_forward(g, logit_scale=30.0) + gen, committed = diffusion_generate( + fwd, tokens, positions, gen_mask, mask_id=MASK_ID, bd_size=4, threshold=0.9, temperature=0.0 + ) + gen = np.asarray(gen) + # Prompt region identical to the input for every row. + assert np.array_equal(gen[:, :prompt_len], np.asarray(tokens)[:, :prompt_len]) + # Prompt positions are never reported as committed. + assert not np.any(np.asarray(committed)[:, :prompt_len]) + + def test_step_cap_force_fill_no_mask_leak(self): + # Cap denoise steps far below what full commitment needs; the post-loop + # force-fill must still leave no mask_id and follow the shifted convention. + prompt_len = 4 + g, tokens, positions, gen_mask = _make_batch(batch_size=1, seq_len=8, prompt_len=prompt_len) + fwd = _make_stub_forward(g, logit_scale=0.05) # below threshold + gen, _ = diffusion_generate( + fwd, + tokens, + positions, + gen_mask, + mask_id=MASK_ID, + bd_size=4, + threshold=0.9, + temperature=0.0, + max_denoise_steps=1, # far fewer than the 4 needed per block + ) + gen = np.asarray(gen)[0] + gnp = np.asarray(g) + assert not np.any(gen == MASK_ID) + for p in range(prompt_len, 8): + assert int(gen[p]) == int(gnp[p - 1]), f"force-filled mismatch at {p}" + + +class TestDiffusionCommit: + """Direct checks on the pure commit primitive.""" + + def test_high_confidence_commits_all_masked(self): + logits = jax.nn.one_hot(jnp.array([1, 2, 3]), VOCAB, dtype=jnp.float32) * 30.0 # [3, V] + mask = jnp.array([True, True, True]) + tok, new_mask = diffusion_commit(logits, mask, threshold=0.9, temperature=0.0) + assert not np.any(np.asarray(new_mask)) # all committed + assert np.asarray(tok).tolist() == [1, 2, 3] + + def test_low_confidence_commits_exactly_one_forced(self): + # Flat-ish logits below threshold -> only the single highest-confidence + # masked position is force-committed; the mask shrinks by exactly one. + logits = jnp.array([[0.10, 0.09, 0.0, 0.0], [0.05, 0.05, 0.05, 0.05]], dtype=jnp.float32) + mask = jnp.array([True, True]) + _, new_mask = diffusion_commit(logits, mask, threshold=0.9, temperature=0.0) + assert int(np.asarray(new_mask).sum()) == 1 + + def test_no_masked_positions_forces_nothing(self): + logits = jnp.zeros((2, VOCAB), dtype=jnp.float32) + mask = jnp.array([False, False]) + _, new_mask = diffusion_commit(logits, mask, threshold=0.9, temperature=0.0) + assert not np.any(np.asarray(new_mask)) # stays empty; nothing forced + + +class TestGenerationIsStopGradient: + """(c) Generation is off the student gradient path.""" + + def test_no_grad_through_rollout_to_forward_params(self): + # forward_fn depends on a differentiable parameter `w`; a scalar built from + # the generated tokens must have zero gradient w.r.t. `w` (argmax/commit is + # non-differentiable), i.e. any loss on the rollout gets no grad from it. + g, tokens, positions, gen_mask = _make_batch(batch_size=1, seq_len=8, prompt_len=4) + g = jnp.asarray(g) + + def scalar_of_w(w): + def fwd(toks, pos): + del toks + return jax.nn.one_hot(g[pos], VOCAB, dtype=jnp.float32) * 30.0 + w + + gen, _ = diffusion_generate(fwd, tokens, positions, gen_mask, mask_id=MASK_ID, bd_size=4, threshold=0.9) + return jnp.sum(gen.astype(jnp.float32)) + + grad = jax.grad(scalar_of_w)(jnp.float32(0.7)) + assert float(grad) == 0.0 + + def test_train_step_wires_rollout_under_stop_gradient(self): + src = _TRAIN_DISTILL.read_text() + # Gated behind the (default-False) opd_on_policy flag. + assert "if self.opd_on_policy:" in src + assert "self._run_on_policy_rollout(batch, student)" in src + # Rollout outputs are placed under stop_gradient. + method = _find_method(ast.parse(src), "MaxTextDistillationTrainer", "_run_on_policy_rollout") + assert method is not None, "_run_on_policy_rollout not defined" + body_src = ast.get_source_segment(src, method) + assert "jax.lax.stop_gradient(gen_tokens)" in body_src + assert "jax.lax.stop_gradient(committed_mask)" in body_src + + +class TestCommittedMaskGatesForwardKL: + """(b) The committed-position mask makes forward-KL nonzero only there.""" + + def _strategy(self, du, vocab, alpha=1.0, temperature=1.0): + return du.CombinedDistillationStrategy( + student_forward_fn=None, + teacher_forward_fn=None, + vocab_size=vocab, + pad_id=0, + temperature=temperature, + alpha=alpha, + ) + + def test_create_labels_mask_is_committed_positions(self): + du = _import_distillation_utils() + strat = self._strategy(du, vocab=6) + targets = jnp.array([[2, 3, 4, 5]]) + committed = jnp.array([[0, 1, 0, 1]]) # committed at positions 1 and 3 + labels = strat.create_labels(targets, targets_segmentation=committed) + mask = np.asarray(jnp.any(labels != 0, axis=-1))[0] + assert mask.tolist() == [False, True, False, True] + + def test_forward_kl_counts_only_committed_positions(self): + du = _import_distillation_utils() + strat = self._strategy(du, vocab=6, alpha=1.0, temperature=1.0) + targets = jnp.array([[2, 3, 4, 5]]) + committed = jnp.array([[0, 1, 0, 1]]) + labels = strat.create_labels(targets, targets_segmentation=committed) + + s_logits = jax.random.normal(jax.random.PRNGKey(0), (1, 4, 6)) + t_logits = jax.random.normal(jax.random.PRNGKey(1), (1, 4, 6)) + s_out = du.DistillationForwardOutput(logits=s_logits) + t_out = du.DistillationForwardOutput(logits=t_logits) + + _, metrics = strat.compute_loss(s_out, t_out, labels, step=None) + kl_sum = float(metrics[du.METRIC_KL_DIV_AT_T][0]) + + def fkl(sl, tl): # forward KL(teacher || student) at one position + return float(jnp.sum(jax.nn.softmax(tl) * (jax.nn.log_softmax(tl) - jax.nn.log_softmax(sl)))) + + hand = fkl(s_logits[0, 1], t_logits[0, 1]) + fkl(s_logits[0, 3], t_logits[0, 3]) + assert abs(kl_sum - hand) < 1e-4 + + def test_loss_invariant_to_non_committed_positions(self): + du = _import_distillation_utils() + strat = self._strategy(du, vocab=6, alpha=1.0, temperature=1.0) + targets = jnp.array([[2, 3, 4, 5]]) + committed = jnp.array([[0, 1, 0, 1]]) + labels = strat.create_labels(targets, targets_segmentation=committed) + + s_logits = jax.random.normal(jax.random.PRNGKey(0), (1, 4, 6)) + t_logits = jax.random.normal(jax.random.PRNGKey(1), (1, 4, 6)) + s_out = du.DistillationForwardOutput(logits=s_logits) + loss_a, _ = strat.compute_loss(s_out, du.DistillationForwardOutput(logits=t_logits), labels, step=None) + # Perturb teacher logits ONLY at non-committed positions 0 and 2. + t2 = t_logits.at[0, 0].add(5.0).at[0, 2].add(-3.0) + loss_b, _ = strat.compute_loss(s_out, du.DistillationForwardOutput(logits=t2), labels, step=None) + assert abs(float(loss_a) - float(loss_b)) < 1e-5 + + +class TestSourceStructure: + """Structural checks that the rollout is on-device and correctly wired.""" + + def test_generate_uses_lax_loops(self): + src = _DIFFUSION_GENERATE.read_text() + assert "jax.lax.while_loop" in src # inner per-block denoise + assert "jax.lax.fori_loop" in src # outer block advance + + def test_train_distill_imports_diffusion_generate(self): + src = _TRAIN_DISTILL.read_text() + assert "import diffusion_generate" in src or "diffusion_generate," in src + assert "diffusion_generate.diffusion_generate(" in src + + +def _find_method(module: ast.Module, cls_name: str, method_name: str): + for node in module.body: + if isinstance(node, ast.ClassDef) and node.name == cls_name: + for sub in node.body: + if isinstance(sub, ast.FunctionDef) and sub.name == method_name: + return sub + return None diff --git a/tests/post_training/unit/rl_diffusion_rollout_test.py b/tests/post_training/unit/rl_diffusion_rollout_test.py new file mode 100644 index 0000000000..3e393f2221 --- /dev/null +++ b/tests/post_training/unit/rl_diffusion_rollout_test.py @@ -0,0 +1,413 @@ +# Copyright 2023–2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU unit tests for block-diffusion RL (GRPO/DAPO) rollout + shared logprob. + +Covers the pure logic of the block-diffusion RL pieces on CPU: + + (a) the shared block-diffusion per-token logprob fn + (``diffusion_generate.diffusion_per_token_logps``): correct shape, the + shifted-logit convention (post-shift logits at ``p`` score token ``p``), + hand-computed values against a stub forward, temperature scaling, and + ``completion_mask`` zeroing; + (b) ``MaxTextDiffusionRollout.generate`` with a stub model produces a valid + completion (no ``mask_id`` leak, EOS truncation, prompt echoed left-padded); + (c) THE P0 INVARIANT: the rollout's old-policy logps + (``MaxTextDiffusionRollout.get_per_token_logps``) and the loss's new-policy + logps (the ``make_diffusion_per_token_logps_fn`` hook) are the SAME function + and return identical values on the same (prompt, completion, stub model). + +Both modules under test are pure/light: ``diffusion_generate.py`` imports only jax +and is loaded standalone by path; ``maxtext_diffusion_rollout.py`` is loaded +standalone with the real (standalone) ``diffusion_generate`` injected and a minimal +``tunix.rl.rollout.base_rollout`` stub, so the whole file runs on a bare jax[cpu] +venv with no MaxText package import or full tunix install. The tunix loss-side +gating (``grpo_loss_fn`` default = AR) is covered in the tunix repo's test. +""" + +import abc +import dataclasses +import importlib.util +import pathlib +import sys +import types +from typing import Any + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from flax import nnx + +pytestmark = [pytest.mark.cpu_only, pytest.mark.post_training] + +_REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +_DIFFUSION_GENERATE = ( + _REPO_ROOT / "src" / "maxtext" / "trainers" / "post_train" / "distillation" / "diffusion_generate.py" +) +_DIFFUSION_ROLLOUT = _REPO_ROOT / "src" / "maxtext" / "integration" / "vllm" / "maxtext_diffusion_rollout.py" +_TRAIN_RL = _REPO_ROOT / "src" / "maxtext" / "trainers" / "post_train" / "rl" / "train_rl.py" + + +def _load_by_path(name, path): + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +# ---- Load diffusion_generate standalone (imports only jax) ------------------ +_DG = _load_by_path("dgr2_diffusion_generate_under_test", _DIFFUSION_GENERATE) +diffusion_per_token_logps = _DG.diffusion_per_token_logps + + +def _install_stub_base_rollout(): + """Minimal ``tunix.rl.rollout.base_rollout`` so the rollout module loads bare.""" + + @dataclasses.dataclass + class RolloutConfig: + max_tokens_to_generate: int = 8 + max_prompt_length: int = 8 + temperature: float = 1.0 + top_p: float | None = 1.0 + top_k: int | None = None + seed: Any = None + + @dataclasses.dataclass + class RolloutOutput: + text: list + logits: Any + tokens: list + left_padded_prompt_tokens: Any + logprobs: Any + + @dataclasses.dataclass + class CacheConfig: + cache_size: int = 0 + num_layers: int = 0 + num_kv_heads: int = 0 + head_dim: int = 0 + + class BaseRollout(abc.ABC): + pass + + base = types.ModuleType("tunix.rl.rollout.base_rollout") + base.RolloutConfig = RolloutConfig + base.RolloutOutput = RolloutOutput + base.CacheConfig = CacheConfig + base.BaseRollout = BaseRollout + for pkg in ("tunix", "tunix.rl", "tunix.rl.rollout"): + if pkg not in sys.modules: + m = types.ModuleType(pkg) + m.__path__ = [] + sys.modules[pkg] = m + sys.modules["tunix.rl.rollout.base_rollout"] = base + sys.modules["tunix.rl.rollout"].base_rollout = base + return base + + +def _install_stub_maxtext_diffusion_generate(): + """Register the standalone ``diffusion_generate`` under its maxtext dotted path.""" + dotted = "maxtext.trainers.post_train.distillation" + parts = dotted.split(".") + for i in range(1, len(parts) + 1): + name = ".".join(parts[:i]) + if name not in sys.modules: + m = types.ModuleType(name) + m.__path__ = [] + sys.modules[name] = m + sys.modules[dotted].diffusion_generate = _DG + sys.modules[dotted + ".diffusion_generate"] = _DG + + +_BASE = _install_stub_base_rollout() +_install_stub_maxtext_diffusion_generate() +_DR = _load_by_path("dgr2_maxtext_diffusion_rollout_under_test", _DIFFUSION_ROLLOUT) +MaxTextDiffusionRollout = _DR.MaxTextDiffusionRollout +make_diffusion_per_token_logps_fn = _DR.make_diffusion_per_token_logps_fn + + +# ---- Test constants and stubs ---------------------------------------------- +VOCAB = 16 +MASK_ID = 15 # outside the normal token range so leaks are detectable +PAD_ID = 0 +EOS_ID = 9 + + +def _target_map(n: int) -> jnp.ndarray: + """Deterministic absolute-position -> token map in [1, VOCAB-3].""" + return jnp.asarray([1 + (p % (VOCAB - 3)) for p in range(n)], dtype=jnp.int32) + + +class _StubModel(nnx.Module): + """Adapter-shaped stub: ``(tokens, positions, cache, attn_mask) -> (logits, None)``. + + Logits peak on ``target[positions]`` (canvas-independent, so argmax/logps are + exact and position-addressed), scaled by ``scale``. Carries one real nnx.Param so + it round-trips through ``nnx.split``/``nnx.merge`` (needed by the loss-side hook). + """ + + def __init__(self, target: jnp.ndarray, scale: float): + self.target = nnx.Variable(jnp.asarray(target, dtype=jnp.int32)) + self.bias = nnx.Param(jnp.zeros((), dtype=jnp.float32)) + self.scale = float(scale) + self.vocab = int(VOCAB) + + def __call__(self, input_tokens, positions, cache=None, attention_mask=None, decoder_segment_ids=None): + del input_tokens, cache, attention_mask, decoder_segment_ids + tgt = self.target[...][positions] # [B, T] + logits = jax.nn.one_hot(tgt, self.vocab, dtype=jnp.float32) * self.scale + self.bias[...] + return logits, None + + +class _StubTokenizer: + """Whitespace tokenizer over integer strings, e.g. "3 4 5" <-> [3, 4, 5].""" + + def encode(self, text: str): + return [int(t) for t in text.split()] if text.strip() else [] + + def decode(self, ids): + return " ".join(str(int(i)) for i in ids) + + def pad_id(self): + return PAD_ID + + def eos_id(self): + return EOS_ID + + +def _make_config(bd_size=4, mask_id=MASK_ID, threshold=0.9, max_denoise_steps=0): + return types.SimpleNamespace( + enable_block_diffusion=True, + bd_size=bd_size, + mask_id=mask_id, + rl=types.SimpleNamespace( + rl_diffusion_threshold=threshold, + rl_diffusion_max_denoise_steps=max_denoise_steps, + ), + ) + + +def _make_rollout(scale=30.0, bd_size=4, temperature=1.0, target_len=64): + model = _StubModel(_target_map(target_len), scale=scale) + rollout_config = _BASE.RolloutConfig(max_tokens_to_generate=4, max_prompt_length=4, temperature=temperature) + rollout = MaxTextDiffusionRollout( + rollout_actor=model, + tokenizer=_StubTokenizer(), + mesh=None, + rollout_config=rollout_config, + maxtext_config=_make_config(bd_size=bd_size), + ) + return rollout, model, rollout_config + + +# ============================================================================ +# (a) shared block-diffusion per-token logprob fn +# ============================================================================ +class TestSharedLogprobFn: + + def _model_apply(self, scale=30.0, target_len=64): + g = _target_map(target_len) + + def model_apply(tokens, positions): + del tokens + return jax.nn.one_hot(g[positions], VOCAB, dtype=jnp.float32) * scale + + return model_apply, g + + def test_shape(self): + fwd, _ = self._model_apply() + prompt = jnp.array([[1, 2, 3, 4], [1, 1, 2, 2]], dtype=jnp.int32) + completion = jnp.array([[3, 4, 5, 6], [2, 3, 4, 5]], dtype=jnp.int32) + logps = diffusion_per_token_logps(fwd, prompt, completion, None, bd_size=4) + assert logps.shape == (2, 4) + + def test_shifted_convention_matches_hand_computed(self): + # positions = 0..7 (all non-pad). Post-shift logits at full position p score + # token p from logits[p-1], which peaks on g[p-1]. So a completion token equal + # to g[p-1] gets logp ~ 0; anything else ~ -scale. + scale = 30.0 + fwd, g = self._model_apply(scale=scale) + prompt = jnp.array([[1, 2, 3, 4]], dtype=jnp.int32) # full positions 0..3 + aligned = jnp.array([[int(g[3]), int(g[4]), int(g[5]), int(g[6])]], dtype=jnp.int32) + logps = np.asarray(diffusion_per_token_logps(fwd, prompt, aligned, None, bd_size=4)) + assert np.allclose(logps, 0.0, atol=1e-4), logps + + # Break one token -> that position drops to ~ -scale, others stay ~ 0. + broken = aligned.at[0, 2].set((int(g[5]) + 1) % (VOCAB - 3) + 1) + logps_b = np.asarray(diffusion_per_token_logps(fwd, prompt, broken, None, bd_size=4)) + assert np.allclose(logps_b[0, [0, 1, 3]], 0.0, atol=1e-4) + assert logps_b[0, 2] < -1.0 + + def test_matches_manual_log_softmax(self): + # Small vocab, moderate scale so the log-softmax is non-degenerate; compare to + # a fully manual computation of the same shifted-logit log-prob. + scale = 1.3 + fwd, g = self._model_apply(scale=scale) + prompt = jnp.array([[2, 3, 4, 5]], dtype=jnp.int32) + completion = jnp.array([[1, 2, 3, 4]], dtype=jnp.int32) + logps = np.asarray(diffusion_per_token_logps(fwd, prompt, completion, None, bd_size=4)) + full = np.concatenate([np.asarray(prompt)[0], np.asarray(completion)[0]]) # [8] + for i in range(4): + p = 4 + i # full position of completion token i + logits_row = np.asarray(jax.nn.one_hot(int(g[p - 1]), VOCAB)) * scale # logits[p-1] + manual = logits_row[full[p]] - np.log(np.sum(np.exp(logits_row))) + assert abs(logps[0, i] - manual) < 1e-4, (i, logps[0, i], manual) + + def test_temperature_scales_logits(self): + fwd, g = self._model_apply(scale=2.0) + prompt = jnp.array([[2, 3, 4, 5]], dtype=jnp.int32) + completion = jnp.array([[1, 2, 3, 4]], dtype=jnp.int32) + a = np.asarray(diffusion_per_token_logps(fwd, prompt, completion, None, bd_size=4, temperature=1.0)) + b = np.asarray(diffusion_per_token_logps(fwd, prompt, completion, None, bd_size=4, temperature=2.0)) + assert not np.allclose(a, b) # temperature changes the distribution + + def test_completion_mask_zeroes_masked_positions(self): + fwd, _ = self._model_apply() + prompt = jnp.array([[1, 2, 3, 4]], dtype=jnp.int32) + completion = jnp.array([[3, 4, 5, 6]], dtype=jnp.int32) + mask = jnp.array([[1, 1, 0, 0]], dtype=jnp.int32) + logps = np.asarray(diffusion_per_token_logps(fwd, prompt, completion, mask, bd_size=4)) + assert np.all(logps[0, 2:] == 0.0) + + def test_requires_divisible_length(self): + fwd, _ = self._model_apply() + prompt = jnp.array([[1, 2, 3]], dtype=jnp.int32) # 3 + 4 = 7, not divisible by 4 + completion = jnp.array([[3, 4, 5, 6]], dtype=jnp.int32) + with pytest.raises(ValueError): + diffusion_per_token_logps(fwd, prompt, completion, None, bd_size=4) + + +# ============================================================================ +# (b) MaxTextDiffusionRollout.generate +# ============================================================================ +class TestDiffusionRolloutGenerate: + + def test_generate_produces_valid_completion(self): + rollout, _, cfg = _make_rollout(scale=30.0, bd_size=4) + out = rollout.generate(["1 2 3 4", "2 3"], cfg) + assert len(out.text) == 2 + assert len(out.tokens) == 2 + # Prompt echoed, left-padded to max_prompt_length=4. + assert out.left_padded_prompt_tokens.shape == (2, 4) + assert out.left_padded_prompt_tokens[0].tolist() == [1, 2, 3, 4] + assert out.left_padded_prompt_tokens[1].tolist() == [PAD_ID, PAD_ID, 2, 3] + for toks in out.tokens: + toks = np.asarray(toks) + assert toks.size >= 1 + assert not np.any(toks == MASK_ID) # no mask token leaks into the output + # EOS-truncated: EOS may appear only as the final token, if at all. + if np.any(toks == EOS_ID): + assert int(toks[-1]) == EOS_ID + assert not np.any(toks[:-1] == EOS_ID) + + def test_requires_block_diffusion_enabled(self): + model = _StubModel(_target_map(64), scale=30.0) + cfg = _make_config() + cfg.enable_block_diffusion = False + with pytest.raises(ValueError): + MaxTextDiffusionRollout( + rollout_actor=model, + tokenizer=_StubTokenizer(), + mesh=None, + rollout_config=_BASE.RolloutConfig(temperature=1.0), + maxtext_config=cfg, + ) + + +# ============================================================================ +# (c) THE P0 INVARIANT +# ============================================================================ +class TestP0Invariant: + + def test_rollout_and_loss_hook_are_same_fn(self): + # Both sites funnel through the SAME single-source-of-truth object. + assert _DR.diffusion_generate.diffusion_per_token_logps is diffusion_per_token_logps + hook = make_diffusion_per_token_logps_fn(bd_size=4, mask_id=MASK_ID) + assert callable(hook) + + def test_rollout_logps_equals_loss_hook_logps(self): + """Old-policy (rollout) logps == new-policy (loss hook) logps, identically.""" + rollout, model, _ = _make_rollout(scale=3.0, bd_size=4, temperature=1.0) + prompt = jnp.array([[1, 2, 3, 4], [1, 1, 2, 2]], dtype=jnp.int32) + completion = jnp.array([[3, 4, 5, 6], [2, 3, 4, 5]], dtype=jnp.int32) + completion_mask = jnp.array([[1, 1, 1, 0], [1, 1, 0, 0]], dtype=jnp.int32) + + # Site 1: rollout old-policy logps. + old_logps = np.asarray(rollout.get_per_token_logps(prompt, completion, completion_mask)) + + # Site 2: loss-side new-policy logps via the pluggable hook (drop-in for + # common.compute_per_token_logps: called with graphdef, state, ...). + hook = make_diffusion_per_token_logps_fn(bd_size=4, mask_id=MASK_ID) + graphdef, state = nnx.split(model) + new_logps = np.asarray( + hook( + graphdef, + state, + prompt_tokens=prompt, + completion_tokens=completion, + pad_id=PAD_ID, + eos_id=EOS_ID, + completion_mask=completion_mask, + stop_gradient=False, + temperature=1.0, + ) + ) + + assert old_logps.shape == new_logps.shape == (2, 4) + np.testing.assert_array_equal(old_logps, new_logps) + + def test_consistency_holds_under_temperature(self): + rollout, model, _ = _make_rollout(scale=2.5, bd_size=4, temperature=0.7) + prompt = jnp.array([[5, 6, 7, 8]], dtype=jnp.int32) + completion = jnp.array([[1, 2, 3, 4]], dtype=jnp.int32) + old_logps = np.asarray(rollout.get_per_token_logps(prompt, completion)) + hook = make_diffusion_per_token_logps_fn(bd_size=4, mask_id=MASK_ID) + graphdef, state = nnx.split(model) + new_logps = np.asarray( + hook( + graphdef, + state, + prompt_tokens=prompt, + completion_tokens=completion, + pad_id=PAD_ID, + eos_id=EOS_ID, + temperature=0.7, + ) + ) + np.testing.assert_array_equal(old_logps, new_logps) + + +# ============================================================================ +# Source-structure checks (train_rl wiring) +# ============================================================================ +class TestTrainRlWiring: + + def test_train_rl_gates_diffusion_rollout(self): + src = _TRAIN_RL.read_text() + assert "rl_diffusion_rollout" in src + assert "MaxTextDiffusionRollout" in src + assert "make_diffusion_per_token_logps_fn" in src + # Diffusion rollout is gated and requires block diffusion. + assert "use_diffusion_rollout" in src + assert "enable_block_diffusion" in src + # The learner receives the shared logprob hook. + assert "per_token_logps_fn=diffusion_per_token_logps_fn" in src + + def test_diffusion_rollout_is_not_a_vllm_subclass(self): + # Fresh JAX rollout (BaseRollout), NOT a VllmRollout subclass. + src = _DIFFUSION_ROLLOUT.read_text() + assert "class MaxTextDiffusionRollout(base_rollout.BaseRollout)" in src + assert "vllm_rollout.VllmRollout" not in src # does not subclass the vLLM rollout + assert "diffusion_generate.diffusion_generate(" in src diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 2841933aa9..ef8d08e8aa 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -42,6 +42,7 @@ from maxtext.layers import attention_op from maxtext.layers.attention_op import ( AttentionOp, + BlockDiffusionMask, ChunkedCausalMask, _generate_chunk_attention_mask, _make_bidirectional_block_mask, @@ -328,6 +329,106 @@ def test_value_error_on_zero_chunk_size(self): _generate_chunk_attention_mask(mask_shape=(4, 4), chunk_size=0) +class BlockDiffusionMaskTest(unittest.TestCase): + """Tests for block-diffusion attention. + + Block-diffusion attention is bidirectional within a block of `bd` tokens and + block-causal across blocks: query position i attends key position j iff + ``(i // bd) >= (j // bd)`` (full attention within the same block and to all + earlier blocks, nothing in later blocks). + """ + + def _expected_block_diffusion(self, q_len, kv_len, bd): + expected = np.zeros((q_len, kv_len), dtype=np.bool_) + for r in range(q_len): + for c in range(kv_len): + if r // bd >= c // bd: + expected[r, c] = True + return expected + + def test_splash_mask_matches_expected(self): + """The splash BlockDiffusionMask yields True iff (i // bd) >= (j // bd).""" + seq_len = 12 + bd = 4 + mask = BlockDiffusionMask(shape=(seq_len, seq_len), bd_size=bd) + expected = self._expected_block_diffusion(seq_len, seq_len, bd) + np.testing.assert_array_equal(mask[:, :], expected) + + def test_bidirectional_within_block_and_block_causal(self): + """Full within-block (incl. future positions), full earlier blocks, no later blocks.""" + seq_len = 12 + bd = 4 + mask = np.asarray(BlockDiffusionMask(shape=(seq_len, seq_len), bd_size=bd)[:, :]) + # Query 1 (block 0) attends future position 3 in the same block (bidirectional). + self.assertTrue(mask[1, 3]) + # Query 0 attends only its own block (positions 0-3). + np.testing.assert_array_equal(mask[0], np.array([1] * 4 + [0] * 8, dtype=bool)) + # Query 4 (block 1) attends all of blocks 0 and 1, none of block 2. + np.testing.assert_array_equal(mask[4], np.array([1] * 8 + [0] * 4, dtype=bool)) + # The last query attends everything. + np.testing.assert_array_equal(mask[seq_len - 1], np.ones(seq_len, dtype=bool)) + + def test_non_square_shape(self): + """Works with different query and key sequence lengths.""" + q_len, kv_len, bd = 6, 8, 2 + mask = BlockDiffusionMask(shape=(q_len, kv_len), bd_size=bd) + expected = self._expected_block_diffusion(q_len, kv_len, bd) + np.testing.assert_array_equal(mask[:, :], expected) + + def test_mask_function_matches_dense_path(self): + """(a) dense generate_attention_mask and (b) BlockDiffusionMask.mask_function agree.""" + seq_len = 12 + bd = 4 + expected = self._expected_block_diffusion(seq_len, seq_len, bd) + + # (b) splash mask_function on broadcastable index grids. + q_ids = np.arange(seq_len)[:, None] + kv_ids = np.arange(seq_len)[None, :] + fn_mask = np.asarray(BlockDiffusionMask(shape=(seq_len, seq_len), bd_size=bd).mask_function(q_ids, kv_ids)) + np.testing.assert_array_equal(fn_mask, expected) + + # (a) dense path via AttentionOp.generate_attention_mask (dot_product kernel). + config = types.SimpleNamespace(context_parallel_load_balance=False, context_sharding="context") + mesh = types.SimpleNamespace(shape={}) + op = AttentionOp( + config=config, + num_query_heads=1, + num_kv_heads=1, + max_target_length=seq_len, + mesh=mesh, + attention_kernel="dot_product", + attention_type=AttentionType.BLOCK_DIFFUSION, + bd_size=bd, + ) + query = jnp.zeros((1, seq_len, 1, 8)) + key = jnp.zeros((1, seq_len, 1, 8)) + decoder_segment_ids = jnp.ones((1, seq_len), dtype=jnp.int32) + dense = op.generate_attention_mask(query, key, decoder_segment_ids, MODEL_MODE_TRAIN) + dense_allow = np.asarray(dense == 0.0)[0, 0, 0] + np.testing.assert_array_equal(dense_allow, expected) + # Dense path and splash mask_function agree. + np.testing.assert_array_equal(dense_allow, fn_mask) + # And it is intentionally NOT a plain causal mask (proves bidirectionality within a block). + causal = np.tril(np.ones((seq_len, seq_len), dtype=np.bool_)) + self.assertFalse(np.array_equal(dense_allow, causal)) + + def test_equality_and_hash(self): + """__eq__/__hash__ depend on shape, bd_size, and q_sequence.""" + a = BlockDiffusionMask(shape=(8, 8), bd_size=4) + b = BlockDiffusionMask(shape=(8, 8), bd_size=4) + c = BlockDiffusionMask(shape=(8, 8), bd_size=2) + self.assertEqual(a, b) + self.assertEqual(hash(a), hash(b)) + self.assertNotEqual(a, c) + + def test_value_error_on_non_positive_bd_size(self): + """A ValueError is raised for bd_size <= 0.""" + with self.assertRaises(ValueError): + BlockDiffusionMask(shape=(4, 4), bd_size=0) + with self.assertRaises(ValueError): + BlockDiffusionMask(shape=(4, 4), bd_size=-3) + + class LoadBalancedMaskTest(unittest.TestCase): """Tests for load-balanced Splash masks.""" diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index 4d61b4caaf..4edd75ec95 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -118,6 +118,80 @@ def test_load_balanced_chunk_context_parallel_config(self): self.assertEqual(config.attention_type, "chunk") self.assertTrue(config.context_parallel_load_balance) + def test_block_diffusion_valid_config(self): + """A valid block-diffusion config loads and exposes the new fields.""" + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "steps=1", + "attention_type=block_diffusion", + "enable_block_diffusion=True", + "bd_size=32", + "max_target_length=2048", + "vocab_size=151936", + "mask_id=151669", + "hardware=tpu", + "packing=False", + "dataset_type=synthetic", + "skip_jax_distributed_system=True", + ] + mock_devices = [unittest.mock.MagicMock(slice_index=0) for _ in range(8)] + with unittest.mock.patch("jax.devices", return_value=mock_devices): + config = pyconfig.initialize(argv) + + self.assertTrue(config.enable_block_diffusion) + self.assertEqual(config.attention_type, "block_diffusion") + self.assertEqual(config.bd_size, 32) + self.assertEqual(config.mask_id, 151669) + + def test_block_diffusion_invalid_configs_raise(self): + """The block-diffusion validator rejects bad bd_size, mask_id, and non-divisible length.""" + base = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "steps=1", + "enable_block_diffusion=True", + "vocab_size=151936", + "hardware=tpu", + "packing=False", + "dataset_type=synthetic", + "skip_jax_distributed_system=True", + ] + bad_overrides = [ + ["bd_size=0", "max_target_length=2048", "mask_id=100"], # bd_size must be > 0 + ["bd_size=32", "max_target_length=2048", "mask_id=999999"], # mask_id >= vocab_size + ["bd_size=7", "max_target_length=2048", "mask_id=100"], # max_target_length % bd_size != 0 + ] + mock_devices = [unittest.mock.MagicMock(slice_index=0) for _ in range(8)] + for overrides in bad_overrides: + with self.assertRaises(pydantic.ValidationError): + with unittest.mock.patch("jax.devices", return_value=mock_devices): + pyconfig.initialize(base + overrides) + + def test_block_diffusion_disabled_skips_validation(self): + """When disabled, the (deliberately out-of-range) block-diffusion defaults do not error.""" + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "steps=1", + "enable_block_diffusion=False", + "hardware=tpu", + "packing=False", + "dataset_type=synthetic", + "skip_jax_distributed_system=True", + ] + mock_devices = [unittest.mock.MagicMock(slice_index=0) for _ in range(8)] + with unittest.mock.patch("jax.devices", return_value=mock_devices): + config = pyconfig.initialize(argv) + + # The default mask_id (151669) exceeds the default vocab_size (32000) but is + # not validated because block diffusion is disabled. + self.assertFalse(config.enable_block_diffusion) + self.assertGreater(config.mask_id, config.vocab_size) + @unittest.mock.patch.dict(os.environ, {pyconfig.yaml_key_to_env_key("steps"): "123"}) def test_env_override(self): """Tests that environment variables override YAML values.""" diff --git a/tests/unit/input_pipeline_utils_test.py b/tests/unit/input_pipeline_utils_test.py index 0d5adae4a9..e014c57d17 100644 --- a/tests/unit/input_pipeline_utils_test.py +++ b/tests/unit/input_pipeline_utils_test.py @@ -17,7 +17,13 @@ import pytest import unittest -from maxtext.input_pipeline.input_pipeline_utils import compute_file_sharding +import numpy as np + +from maxtext.input_pipeline.input_pipeline_utils import ( + BlockDiffusionMasking, + ShiftData, + compute_file_sharding, +) @pytest.mark.cpu_only @@ -109,5 +115,137 @@ def test_no_row_shard_when_only_one_reader(self): self.assertIsNone(row_shard) +_PAD_ID = 0 +_MASK_ID = 99 +_BD_SIZE = 4 + + +def _make_sft_batch(): + """Synthetic post-PadOrTrim, post-Batch, pre-shift SFT batch (completion-only). + + Two examples, seq_len 8 (== 2 blocks of bd_size 4). Prompt and padding carry + targets_segmentation == 0; the completion span carries targets_segmentation == 1. + All token ids are nonzero (so pad_id 0 marks only prompt-mask/padding) and none + equal _MASK_ID. + """ + # ex0: prompt [11,12], response [21,22,23,24], pad [0,0] + # ex1: prompt [31,32,33], response [41,42,43], pad [0,0] + inputs = np.array( + [ + [11, 12, 21, 22, 23, 24, 0, 0], + [31, 32, 33, 41, 42, 43, 0, 0], + ], + dtype=np.int32, + ) + # completion-only targets: prompt + padding zeroed, response = clean tokens (aligned). + targets = np.array( + [ + [0, 0, 21, 22, 23, 24, 0, 0], + [0, 0, 0, 41, 42, 43, 0, 0], + ], + dtype=np.int32, + ) + inputs_segmentation = (inputs != _PAD_ID).astype(np.int32) + targets_segmentation = (targets != _PAD_ID).astype(np.int32) + positions = np.broadcast_to(np.arange(inputs.shape[1], dtype=np.int32), inputs.shape).copy() + return { + "inputs": inputs, + "targets": targets, + "inputs_segmentation": inputs_segmentation, + "targets_segmentation": targets_segmentation, + "inputs_position": positions.copy(), + "targets_position": positions.copy(), + } + + +@pytest.mark.cpu_only +class BlockDiffusionMaskingTest(unittest.TestCase): + """Tests for the masked-diffusion (CFT) data transform for block-diffusion SFT.""" + + def _apply(self, seed): + element = _make_sft_batch() + transform = BlockDiffusionMasking(bd_size=_BD_SIZE, mask_id=_MASK_ID) + return transform.random_map(element, np.random.default_rng(seed)) + + def test_only_response_positions_are_masked(self): + """(a) Prompt and padding are never masked; masks are a subset of the response span.""" + clean = _make_sft_batch() + response = clean["targets_segmentation"] != 0 + for seed in range(50): + out = self._apply(seed) + masked = out["inputs"] == _MASK_ID + # Every masked position is a response position. + self.assertTrue(np.all(masked <= response), f"seed {seed}: masked outside response span") + # Non-response inputs (prompt + padding) are byte-for-byte unchanged. + np.testing.assert_array_equal(out["inputs"][~response], clean["inputs"][~response]) + + def test_masked_inputs_equal_mask_id(self): + """(b) Masked input positions hold mask_id, and some masking actually happens.""" + out = self._apply(seed=0) + masked = out["targets_segmentation"] != 0 + self.assertGreater(int(masked.sum()), 0, "expected at least one masked position") + self.assertTrue(np.all(out["inputs"][masked] == _MASK_ID)) + + def test_targets_are_clean_and_aligned(self): + """(c) Targets are the clean tokens, aligned (identical to input, not next-token shifted).""" + clean = _make_sft_batch() + for seed in range(10): + out = self._apply(seed) + np.testing.assert_array_equal(out["targets"], clean["targets"]) + # And explicitly NOT the AR next-token-shifted targets. + shifted = np.array([row[1:].tolist() + [_PAD_ID] for row in clean["targets"]], dtype=np.int32) + self.assertFalse(np.array_equal(out["targets"], shifted)) + + def test_targets_segmentation_is_one_exactly_at_masked(self): + """(d) targets_segmentation == 1 exactly at masked positions, 0 everywhere else.""" + out = self._apply(seed=0) + masked = (out["inputs"] == _MASK_ID).astype(np.int32) + np.testing.assert_array_equal(out["targets_segmentation"], masked) + clean = _make_sft_batch() + non_response = clean["targets_segmentation"] == 0 + self.assertTrue(np.all(out["targets_segmentation"][non_response] == 0)) + + def test_inputs_segmentation_covers_real_tokens(self): + """(e) inputs_segmentation is unchanged and covers all real prompt+response tokens.""" + clean = _make_sft_batch() + out = self._apply(seed=0) + np.testing.assert_array_equal(out["inputs_segmentation"], clean["inputs_segmentation"]) + # Masked positions are still attended to (segmentation stays 1 there). + masked = out["inputs"] == _MASK_ID + self.assertTrue(np.all(out["inputs_segmentation"][masked] == 1)) + + def test_determinism_for_fixed_seed(self): + """(f) A fixed RNG seed yields identical corruption; different seeds differ.""" + out_a = self._apply(seed=123) + out_b = self._apply(seed=123) + for key, value_a in out_a.items(): + np.testing.assert_array_equal(value_a, out_b[key], err_msg=f"non-deterministic key {key}") + out_c = self._apply(seed=124) + self.assertFalse(np.array_equal(out_a["inputs"], out_c["inputs"])) + + def test_flag_off_ar_path_shifts_targets(self): + """(g) The AR path (ShiftData, used when enable_block_diffusion=False) shifts targets; + + block diffusion keeps them aligned. This confirms the two pipeline branches are + genuinely different objectives and that leaving the flag off preserves the AR shift. + """ + ar = ShiftData(ignored_ids=[_PAD_ID]).map(_make_sft_batch()) + clean = _make_sft_batch() + # ShiftData shifts targets left by one (next-token prediction). + np.testing.assert_array_equal(ar["targets"][:, :-1], clean["targets"][:, 1:]) + # Block diffusion leaves targets aligned (no shift). + bd = self._apply(seed=0) + np.testing.assert_array_equal(bd["targets"], clean["targets"]) + self.assertFalse(np.array_equal(ar["targets"], bd["targets"])) + + def test_invalid_bd_size_raises(self): + """bd_size must be positive and must divide the sequence length.""" + with self.assertRaises(ValueError): + BlockDiffusionMasking(bd_size=0, mask_id=_MASK_ID) + with self.assertRaises(ValueError): + # seq_len 8 is not divisible by bd_size 3. + BlockDiffusionMasking(bd_size=3, mask_id=_MASK_ID).random_map(_make_sft_batch(), np.random.default_rng(0)) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/model_test.py b/tests/unit/model_test.py index 41cad97198..9fd5f0ea6b 100644 --- a/tests/unit/model_test.py +++ b/tests/unit/model_test.py @@ -197,5 +197,94 @@ def test_train_vs_prefill_and_autoregress(self): ) +class Qwen3BlockDiffusionForwardTest(unittest.TestCase): + """End-to-end CPU forward of a tiny Qwen3 with block-diffusion attention. + + Block diffusion is weight-preserving: only the attention mask changes, not the + weights. This test applies the *same* parameters through both a block-diffusion + Qwen3 and a causal (global) Qwen3 and asserts (1) the parameter tree is shared + by both (weight-preserving) and (2) the logits differ, proving the + `enable_block_diffusion` flag is wired into the Qwen3 forward pass. + """ + + def _init_qwen3_config(self, enable_block_diffusion): + """Builds a tiny, CPU-friendly Qwen3 config (dot_product attention).""" + return pyconfig.initialize( + [sys.argv[0], get_test_config_path()], + model_name="qwen3-0.6b", + override_model_config=True, + per_device_batch_size=1.0, + run_name="test", + enable_checkpointing=False, + base_num_decoder_layers=2, + attention="dot_product", + max_target_length=16, + base_emb_dim=64, + base_num_query_heads=2, + base_num_kv_heads=1, + head_dim=32, + vocab_size=256, + max_prefill_predict_length=4, + dataset_type="synthetic", + enable_block_diffusion=enable_block_diffusion, + bd_size=4, + mask_id=100, + ) + + def test_block_diffusion_changes_forward_and_is_weight_preserving(self): + rng = jax.random.PRNGKey(0) + cfg_bd = self._init_qwen3_config(enable_block_diffusion=True) + cfg_causal = self._init_qwen3_config(enable_block_diffusion=False) + self.assertTrue(cfg_bd.enable_block_diffusion) + self.assertFalse(cfg_causal.enable_block_diffusion) + + mesh = Mesh(maxtext_utils.create_device_mesh(cfg_bd), cfg_bd.mesh_axes) + model_bd = models.transformer_as_linen(config=cfg_bd, mesh=mesh, quant=None, model_mode=MODEL_MODE_TRAIN) + model_causal = models.transformer_as_linen(config=cfg_causal, mesh=mesh, quant=None, model_mode=MODEL_MODE_TRAIN) + + shape = (cfg_bd.global_batch_size_to_train_on, cfg_bd.max_target_length) + ids = jax.random.randint(rng, shape, 0, cfg_bd.vocab_size) + decoder_segment_ids = jnp.zeros(shape) + DECODING_ACTIVE_SEQUENCE_INDICATOR + decoder_positions = jnp.stack( + [jnp.arange(cfg_bd.max_target_length, dtype=jnp.int32) for _ in range(cfg_bd.global_batch_size_to_train_on)] + ) + + # Weight-preserving: initialize once, then apply the same params through both models. + params = model_causal.init( + {"params": rng, "aqt": rng, "dropout": rng}, + ids, + decoder_positions, + decoder_segment_ids, + enable_dropout=False, + ) + out_causal = np.asarray( + model_causal.apply( + params, + ids, + decoder_positions, + decoder_segment_ids, + enable_dropout=False, + model_mode=MODEL_MODE_TRAIN, + rngs={"aqt": rng}, + ) + ) + out_bd = np.asarray( + model_bd.apply( + params, + ids, + decoder_positions, + decoder_segment_ids, + enable_dropout=False, + model_mode=MODEL_MODE_TRAIN, + rngs={"aqt": rng}, + ) + ) + + self.assertEqual(out_bd.shape, (shape[0], shape[1], cfg_bd.vocab_size)) + # The block-diffusion mask (bidirectional within a block) differs from the + # causal mask, so identical weights must yield different logits. + self.assertFalse(np.allclose(out_causal, out_bd)) + + if __name__ == "__main__": unittest.main()