diff --git a/src/maxtext/common/common_types.py b/src/maxtext/common/common_types.py index 9dd4aff3dc..8ed18bfd12 100644 --- a/src/maxtext/common/common_types.py +++ b/src/maxtext/common/common_types.py @@ -135,6 +135,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 55925ffd58..bd77734d16 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -380,12 +380,13 @@ 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_block_size: 32 # Size of each bidirectional block when attention_type='block_diffusion' attn_logits_soft_cap: 0.0 final_logits_soft_cap: 0.0 z_loss_multiplier: 0.0 @@ -790,6 +791,11 @@ olmo_apply_ngram_filter: true # mask instances with repetitive n-grams (OLMo-cor # Training loop steps: 150_001 # If set to -1 then will inherit value from learning_rate_schedule_steps log_period: 100 # The frequency of Tensorboard flush, gcs metrics writing, and managed profiler metrics updating. +training_objective: 'causal_lm' # Supported objectives: causal_lm, block_diffusion +block_diffusion_mask_id: -1 # Tokenizer mask-token id; required for training_objective='block_diffusion' +block_diffusion_min_noise: 0.001 # Minimum per-block corruption probability for block-diffusion training +block_diffusion_logit_alignment: 'same_position' # Supported alignments: same_position, shifted +block_diffusion_canvas_policy: 'all_masked' # Supported canvases: all_masked, seed_and_mask jax_distributed_initialization_timeout: 300 # This is the default timeout in https://github.com/jax-ml/jax/blob/main/jax/_src/distributed.py # Note there are two separate initializations - the jax coordination service (aka jax.distributed.initialize) and the backend (e.g. PjRT), the timeout above refers @@ -1315,6 +1321,12 @@ engram_kernel_size: 4 engram_seed: 0 ##### Distillation parameters +distill_data_source: "dataset" # dataset or student_rollout for diffusion on-policy distillation +distill_rollout_algorithm: "low_confidence" +distill_rollout_confidence_threshold: 0.9 +distill_rollout_temperature: 1.0 +distill_rollout_max_denoise_steps: -1 # -1 uses block_diffusion_block_size +distill_rollout_stop_token_ids: [] # Empty uses the tokenizer EOS; set model-specific EOT IDs explicitly distill_alpha: 0.5 distill_temperature: 1.0 # distill_beta is used for cosine similarity loss between intermediate activataitions of out_proj in teacher/student models. diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 03640ec758..4fcdc5f289 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -572,7 +572,7 @@ class Attention(BaseModel): "autoselected", description="The attention algorithm to use (dot_product, flash, cudnn_flash_te, vllm_rpa, vllm_batched_rpa, 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( @@ -619,6 +619,10 @@ 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.") + block_diffusion_block_size: PositiveInt = Field( + 32, + description="The size of each bidirectional block when attention_type is 'block_diffusion'.", + ) attn_logits_soft_cap: None | NonNegativeFloat = Field( None, description="Soft-cap value for attention logits. None means no cap." ) @@ -1507,6 +1511,35 @@ class Distillation(BaseModel): description="GCS or local path to the pre-generated ArrayRecord teacher data.", ) + distill_data_source: Literal["dataset", "student_rollout"] = Field( + "dataset", + description="Whether distillation consumes dataset targets or fresh rollouts from the current student.", + ) + distill_rollout_algorithm: Literal["low_confidence"] = Field( + "low_confidence", + description="Block-diffusion rollout algorithm used for on-policy distillation.", + ) + distill_rollout_confidence_threshold: float = Field( + 0.9, + ge=0.0, + le=1.0, + description="Confidence threshold for parallel token commits during diffusion rollout.", + ) + distill_rollout_temperature: float = Field( + 1.0, + gt=0.0, + description="Temperature used to compute rollout token confidence.", + ) + distill_rollout_max_denoise_steps: int = Field( + -1, + ge=-1, + description="Maximum denoising iterations per block; -1 uses block_diffusion_block_size.", + ) + distill_rollout_stop_token_ids: list[int] = Field( + default_factory=list, + description="Generated token IDs that terminate the OPD completion; empty uses the tokenizer EOS ID.", + ) + # --- Loss Params --- distill_alpha: float = Field(0.5, description="Weight for the distillation loss component.") distill_temperature: float = Field(1.0, description="Temperature for distillation softening.") @@ -1579,6 +1612,28 @@ class Distillation(BaseModel): class TrainingLoop(BaseModel): """Configuration for the main training loop, evaluation, and reproducibility.""" + training_objective: Literal["causal_lm", "block_diffusion"] = Field( + "causal_lm", + description="The token-prediction objective used to prepare targets and compute loss.", + ) + block_diffusion_mask_id: int = Field( + -1, + description="The tokenizer mask-token id required by the block-diffusion training objective.", + ) + block_diffusion_min_noise: float = Field( + 1.0e-3, + gt=0.0, + le=1.0, + description="The minimum corruption probability sampled independently for each block.", + ) + block_diffusion_logit_alignment: Literal["same_position", "shifted"] = Field( + "same_position", + description="How model logits align to clean target-token positions.", + ) + block_diffusion_canvas_policy: Literal["all_masked", "seed_and_mask"] = Field( + "all_masked", + description="Whether every block is fully maskable or begins with a clean anchor token.", + ) steps: int = Field( 150_001, ge=-1, @@ -3283,6 +3338,44 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de not isinstance(self.sliding_window_size, int) or self.sliding_window_size <= 0 ): raise ValueError("`sliding_window_size` must be an integer > 0 for 'local_sliding' attention.") + if self.attention_type == AttentionType.BLOCK_DIFFUSION.value: + if self.attention not in ("autoselected", "dot_product", "flash"): + raise ValueError("Block-diffusion attention is supported only by dot_product attention and TPU Splash attention.") + if self.attention in ("autoselected", "flash") and self.hardware != "tpu": + raise ValueError( + "Block-diffusion attention with attention='autoselected' or attention='flash' requires hardware='tpu'; " + "use attention='dot_product' on other hardware." + ) + if self.training_objective == "block_diffusion": + if self.attention_type != AttentionType.BLOCK_DIFFUSION.value: + raise ValueError("`training_objective='block_diffusion'` requires `attention_type='block_diffusion'`.") + if self.block_diffusion_mask_id < 0 or self.block_diffusion_mask_id >= self.vocab_size: + raise ValueError( + f"`block_diffusion_mask_id` ({self.block_diffusion_mask_id}) must satisfy " + f"0 <= block_diffusion_mask_id < vocab_size ({self.vocab_size})." + ) + if self.packing: + raise ValueError("`training_objective='block_diffusion'` requires `packing=False`.") + if self.mtp_num_layers > 0: + raise ValueError("`training_objective='block_diffusion'` is not compatible with MTP.") + if self.num_vocab_tiling > 1: + raise ValueError("`training_objective='block_diffusion'` is not compatible with vocabulary tiling.") + if self.dataset_type != "hf": + raise ValueError("`training_objective='block_diffusion'` currently requires `dataset_type='hf'`.") + if self.use_dpo: + raise ValueError("`training_objective='block_diffusion'` is not compatible with DPO.") + if self.use_multimodal or self.use_audio: + raise ValueError("`training_objective='block_diffusion'` currently supports text-only training.") + valid_model_contracts = { + ("same_position", "all_masked"), + ("shifted", "seed_and_mask"), + } + model_contract = (self.block_diffusion_logit_alignment, self.block_diffusion_canvas_policy) + if model_contract not in valid_model_contracts: + raise ValueError( + "Block-diffusion training supports only `same_position/all_masked` or `shifted/seed_and_mask`; " + f"received `{model_contract[0]}/{model_contract[1]}`." + ) if self.quantize_kvcache and not self.kv_quant_axis: raise ValueError("`kv_quant_axis` cannot be empty when quantize_kvcache is True.") if ( diff --git a/src/maxtext/diffusion/__init__.py b/src/maxtext/diffusion/__init__.py new file mode 100644 index 0000000000..2c9fc63d21 --- /dev/null +++ b/src/maxtext/diffusion/__init__.py @@ -0,0 +1,15 @@ +# Copyright 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 +# +# http://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. + +"""Shared block-diffusion model semantics.""" diff --git a/src/maxtext/diffusion/denoise.py b/src/maxtext/diffusion/denoise.py new file mode 100644 index 0000000000..937c146a49 --- /dev/null +++ b/src/maxtext/diffusion/denoise.py @@ -0,0 +1,348 @@ +# Copyright 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 +# +# http://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. + +"""Model-independent block-diffusion rollout state transitions.""" + +from collections.abc import Callable +from typing import NamedTuple + +import jax +import jax.numpy as jnp +import numpy as np + + +class DenoiseTrace(NamedTuple): + """Compact sampled trajectory for replaying diffusion policy scores.""" + + tokens: jax.Array + action_steps: jax.Array + action_logps: jax.Array + + +def _validate_shapes(initial_tokens, positions, validity_mask, completion_mask): + """Checks that all token-level rollout arrays share a batch-major shape.""" + expected_shape = tuple(initial_tokens.shape) + if len(expected_shape) != 2: + raise ValueError(f"initial_tokens must have shape [batch, length]; received {expected_shape}") + for name, value in ( + ("positions", positions), + ("validity_mask", validity_mask), + ("completion_mask", completion_mask), + ): + if tuple(value.shape) != expected_shape: + raise ValueError(f"{name} must match initial_tokens shape; received {tuple(value.shape)} and {expected_shape}") + + +def _concrete_numpy(value): + if isinstance(value, jax.core.Tracer): + return None + if isinstance(value, jax.Array) and not value.is_fully_addressable: + return None + return np.asarray(value) + + +def _per_row_rngs(rng, batch_size): + """Normalizes one or per-row typed/legacy JAX keys to per-row keys.""" + rng = jnp.asarray(rng) + is_typed_key = jax.dtypes.issubdtype(rng.dtype, jax.dtypes.prng_key) + scalar_shape = () if is_typed_key else tuple(jax.random.PRNGKey(0).shape) + if tuple(rng.shape) == scalar_shape: + return jax.vmap(lambda row: jax.random.fold_in(rng, row))(jnp.arange(batch_size, dtype=jnp.int32)) + if tuple(rng.shape) == (batch_size, *scalar_shape): + return rng + raise ValueError( + "rng must be one PRNG key or one key per batch row; " + f"received shape {tuple(rng.shape)} for key shape {scalar_shape}" + ) + + +def _select_row_rngs(row_mask, new_rngs, old_rngs): + selector = row_mask.reshape((row_mask.shape[0],) + (1,) * (old_rngs.ndim - 1)) + return jnp.where(selector, new_rngs, old_rngs) + + +def _validate_logical_positions(positions, validity_mask, completion_mask, *, shifted_seed): + """Checks logical sequence invariants on eager, host-addressable inputs.""" + concrete_positions = _concrete_numpy(positions) + concrete_validity = _concrete_numpy(validity_mask) + concrete_completion = _concrete_numpy(completion_mask) + if concrete_positions is None or concrete_validity is None or concrete_completion is None: + return + concrete_validity = np.asarray(concrete_validity, dtype=bool) + concrete_completion = np.asarray(concrete_completion, dtype=bool) & concrete_validity + sequence_length = concrete_positions.shape[1] + for row in range(concrete_positions.shape[0]): + valid_positions = np.asarray(concrete_positions[row, concrete_validity[row]]) + expected_positions = np.arange(valid_positions.size, dtype=valid_positions.dtype) + if valid_positions.size and not np.array_equal(np.sort(valid_positions), expected_positions): + raise ValueError( + "valid logical positions must be unique and contiguous from zero within the physical sequence length" + ) + if np.any(valid_positions < 0) or np.any(valid_positions >= sequence_length): + raise ValueError("valid logical positions must be nonnegative and smaller than the physical sequence length") + if shifted_seed: + position_zero = concrete_validity[row] & (concrete_positions[row] == 0) + if np.count_nonzero(position_zero) != 1 or np.any(concrete_completion[row] & position_zero): + raise ValueError("shifted seed-and-mask rollout requires exactly one prompt token at logical position zero") + + +def validate_completion_suffix(positions, validity_mask, completion_mask, *, shifted_seed=False): + """Validates the initial OPD rollout scope when eager values are available. + + The first integration intentionally supports one prompt followed by one + completion per sequence. This prevents clean future turns from appearing in + the bidirectional attention block of an earlier generated completion. + """ + _validate_logical_positions(positions, validity_mask, completion_mask, shifted_seed=shifted_seed) + concrete_positions = _concrete_numpy(positions) + concrete_validity = _concrete_numpy(validity_mask) + concrete_completion = _concrete_numpy(completion_mask) + if concrete_positions is None or concrete_validity is None or concrete_completion is None: + return + concrete_validity = np.asarray(concrete_validity, dtype=bool) + concrete_completion = np.asarray(concrete_completion, dtype=bool) & concrete_validity + for row in range(concrete_positions.shape[0]): + valid_indices = np.flatnonzero(concrete_validity[row]) + if valid_indices.size == 0: + continue + ordered_indices = valid_indices[np.argsort(concrete_positions[row, valid_indices])] + ordered_completion = concrete_completion[row, ordered_indices] + completion_indices = np.flatnonzero(ordered_completion) + if completion_indices.size and not np.all(ordered_completion[completion_indices[0] :]): + raise ValueError("diffusion OPD rollout requires completion_mask to be a contiguous suffix") + + +def low_confidence_rollout( + logits_fn: Callable[[jax.Array], jax.Array], + initial_tokens: jax.Array, + positions: jax.Array, + validity_mask: jax.Array, + completion_mask: jax.Array, + *, + block_size: int, + mask_id: int, + logit_alignment: str, + canvas_policy: str, + confidence_threshold: float = 0.9, + temperature: float = 1.0, + max_denoise_steps: int | None = None, + rng: jax.Array | None = None, +) -> DenoiseTrace: + """Samples a completion and records each token's pre-commit action step. + + ``logits_fn`` must return target-aligned logits for the current token canvas. + Each denoising step commits every token at or above the confidence threshold; + if a row has no such token, its highest-confidence unresolved token is forced + to commit. ``action_steps`` compactly identifies the exact masked canvas for + replay: token ``p`` was scored while every token with step >= its step was + still masked. Prompt, padding, and inactive completion tokens use step -1. + """ + _validate_shapes(initial_tokens, positions, validity_mask, completion_mask) + valid_contracts = {("same_position", "all_masked"), ("shifted", "seed_and_mask")} + if (logit_alignment, canvas_policy) not in valid_contracts: + raise ValueError( + "rollout supports only same_position/all_masked or shifted/seed_and_mask; " + f"received {logit_alignment}/{canvas_policy}" + ) + if block_size <= 0: + raise ValueError(f"block_size must be positive; received {block_size}") + if not 0.0 <= confidence_threshold <= 1.0: + raise ValueError(f"confidence_threshold must be in [0, 1]; received {confidence_threshold}") + if temperature <= 0.0: + raise ValueError(f"temperature must be positive; received {temperature}") + if max_denoise_steps is None: + max_denoise_steps = block_size + if max_denoise_steps < block_size: + raise ValueError( + f"max_denoise_steps must be at least block_size ({block_size}) to guarantee completion; " + f"received {max_denoise_steps}" + ) + + shifted_seed = logit_alignment == "shifted" + validate_completion_suffix(positions, validity_mask, completion_mask, shifted_seed=shifted_seed) + validity_mask = jnp.asarray(validity_mask, dtype=jnp.bool_) + completion_mask = jnp.asarray(completion_mask, dtype=jnp.bool_) & validity_mask + positions = jnp.asarray(positions, dtype=jnp.int32) + canvas = jnp.where(completion_mask, jnp.asarray(mask_id, initial_tokens.dtype), initial_tokens) + action_steps = jnp.full(initial_tokens.shape, -1, dtype=jnp.int32) + action_logps = jnp.zeros(initial_tokens.shape, dtype=jnp.float32) + sample_tokens = rng is not None + rng = jax.random.PRNGKey(0) if rng is None else jnp.asarray(rng) + row_rngs = _per_row_rngs(rng, initial_tokens.shape[0]) + block_ids = positions // block_size + num_blocks = (initial_tokens.shape[1] + block_size - 1) // block_size + + def propose(current_canvas, proposal_keys): + logits = logits_fn(current_canvas) + expected_prefix = tuple(current_canvas.shape) + if len(logits.shape) != 3 or tuple(logits.shape[:2]) != expected_prefix: + raise ValueError( + "logits_fn must return [batch, length, vocab] target-aligned logits; " + f"received {tuple(logits.shape)} for canvas {expected_prefix}" + ) + vocab_size = logits.shape[-1] + if vocab_size < 2: + raise ValueError("logits_fn must expose at least two vocabulary entries so the mask token can be excluded") + if not 0 <= mask_id < vocab_size: + raise ValueError(f"mask_id must satisfy 0 <= mask_id < vocab_size ({vocab_size}); received {mask_id}") + scaled_logits = jnp.asarray(logits, dtype=jnp.float32) / temperature + scaled_logits = scaled_logits.at[..., mask_id].set(-jnp.inf) + log_probabilities = jax.nn.log_softmax(scaled_logits, axis=-1) + if sample_tokens: + proposed_tokens = jax.vmap(lambda key, row_logits: jax.random.categorical(key, row_logits, axis=-1))( + proposal_keys, scaled_logits + ).astype(initial_tokens.dtype) + else: + proposed_tokens = jnp.argmax(scaled_logits, axis=-1).astype(initial_tokens.dtype) + proposed_logps = jnp.take_along_axis(log_probabilities, proposed_tokens[..., None], axis=-1)[..., 0] + confidence = jnp.max(jnp.exp(log_probabilities), axis=-1) + return proposed_tokens, confidence, proposed_logps + + def generate_block(block_id, rollout_state): + current_canvas, current_steps, current_logps, global_step, current_rng = rollout_state + in_block = completion_mask & (block_ids == block_id) + + def run_active_block(active_state): + active_canvas, active_steps, active_logps, active_global_step, active_rng = active_state + anchors = in_block & (positions % block_size == 0) if shifted_seed else jnp.zeros_like(in_block) + + def generate_anchors(anchor_state): + anchor_canvas, anchor_steps, anchor_logps, anchor_global_step, anchor_rng = anchor_state + split_keys = jax.vmap(jax.random.split)(anchor_rng) + next_rng = split_keys[:, 0] + proposal_keys = split_keys[:, 1] + anchor_tokens, _, proposed_logps = propose(anchor_canvas, proposal_keys) + anchor_canvas = jnp.where(anchors, anchor_tokens, anchor_canvas) + anchor_steps = jnp.where(anchors, anchor_global_step[:, None], anchor_steps) + anchor_logps = jnp.where(anchors, proposed_logps, anchor_logps) + rows_with_anchors = jnp.any(anchors, axis=1) + anchor_rng = _select_row_rngs(rows_with_anchors, next_rng, anchor_rng) + return ( + anchor_canvas, + anchor_steps, + anchor_logps, + anchor_global_step + rows_with_anchors.astype(jnp.int32), + anchor_rng, + ) + + if shifted_seed: + active_canvas, active_steps, active_logps, active_global_step, active_rng = jax.lax.cond( + jnp.any(anchors), + generate_anchors, + lambda value: value, + (active_canvas, active_steps, active_logps, active_global_step, active_rng), + ) + unresolved = in_block & ~anchors + + def continue_denoising(state): + step, _, _, _, _, _, remaining = state + return (step < max_denoise_steps) & jnp.any(remaining) + + def denoise_step(state): + step, step_canvas, step_ids, step_logps, step_global, step_rng, remaining = state + split_keys = jax.vmap(jax.random.split)(step_rng) + next_rng = split_keys[:, 0] + proposal_keys = split_keys[:, 1] + proposed_tokens, confidence, proposed_logps = propose(step_canvas, proposal_keys) + commits = remaining & (confidence >= confidence_threshold) + row_needs_fallback = jnp.any(remaining, axis=1) & ~jnp.any(commits, axis=1) + fallback_confidence = jnp.max(jnp.where(remaining, confidence, -jnp.inf), axis=1) + tied_for_fallback = remaining & (confidence == fallback_confidence[:, None]) + fallback_positions = jnp.where(tied_for_fallback, positions, positions.shape[1]) + fallback_indices = jnp.argmin(fallback_positions, axis=1) + fallback = jax.nn.one_hot(fallback_indices, remaining.shape[1], dtype=jnp.bool_) + commits |= fallback & row_needs_fallback[:, None] + step_canvas = jnp.where(commits, proposed_tokens, step_canvas) + step_ids = jnp.where(commits, step_global[:, None], step_ids) + step_logps = jnp.where(commits, proposed_logps, step_logps) + rows_with_commits = jnp.any(commits, axis=1) + step_rng = _select_row_rngs(rows_with_commits, next_rng, step_rng) + return ( + step + 1, + step_canvas, + step_ids, + step_logps, + step_global + rows_with_commits.astype(jnp.int32), + step_rng, + remaining & ~commits, + ) + + _, active_canvas, active_steps, active_logps, active_global_step, active_rng, _ = jax.lax.while_loop( + continue_denoising, + denoise_step, + ( + jnp.asarray(0, dtype=jnp.int32), + active_canvas, + active_steps, + active_logps, + active_global_step, + active_rng, + unresolved, + ), + ) + return active_canvas, active_steps, active_logps, active_global_step, active_rng + + return jax.lax.cond( + jnp.any(in_block), + run_active_block, + lambda value: value, + (current_canvas, current_steps, current_logps, global_step, current_rng), + ) + + canvas, action_steps, action_logps, _, _ = jax.lax.fori_loop( + 0, + num_blocks, + generate_block, + ( + canvas, + action_steps, + action_logps, + jnp.zeros((initial_tokens.shape[0],), dtype=jnp.int32), + row_rngs, + ), + ) + return DenoiseTrace(tokens=canvas, action_steps=action_steps, action_logps=action_logps) + + +def low_confidence_generate( + logits_fn: Callable[[jax.Array], jax.Array], + initial_tokens: jax.Array, + positions: jax.Array, + validity_mask: jax.Array, + completion_mask: jax.Array, + *, + block_size: int, + mask_id: int, + logit_alignment: str, + canvas_policy: str, + confidence_threshold: float = 0.9, + temperature: float = 1.0, + max_denoise_steps: int | None = None, +) -> jax.Array: + """Greedily generates tokens while preserving the original public API.""" + + return low_confidence_rollout( + logits_fn, + initial_tokens, + positions, + validity_mask, + completion_mask, + block_size=block_size, + mask_id=mask_id, + logit_alignment=logit_alignment, + canvas_policy=canvas_policy, + confidence_threshold=confidence_threshold, + temperature=temperature, + max_denoise_steps=max_denoise_steps, + ).tokens diff --git a/src/maxtext/diffusion/scoring.py b/src/maxtext/diffusion/scoring.py new file mode 100644 index 0000000000..bf220219dc --- /dev/null +++ b/src/maxtext/diffusion/scoring.py @@ -0,0 +1,43 @@ +# Copyright 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 +# +# http://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. + +"""Target-alignment utilities shared by diffusion training objectives.""" + +import jax.numpy as jnp + + +def align_logits_to_targets(logits, alignment, positions=None, validity_mask=None): + """Aligns model output positions with clean block-diffusion targets.""" + if alignment == "same_position": + return logits + if alignment == "shifted": + if positions is None: + indices = jnp.maximum(jnp.arange(logits.shape[1], dtype=jnp.int32) - 1, 0) + return logits[:, indices, :] + positions = jnp.asarray(positions, dtype=jnp.int32) + if positions.shape != logits.shape[:2]: + raise ValueError(f"positions must match logits [batch, length]; got {positions.shape} and {logits.shape[:2]}") + if validity_mask is None: + validity_mask = jnp.ones_like(positions, dtype=jnp.bool_) + sequence_length = logits.shape[1] + array_indices = jnp.broadcast_to(jnp.arange(sequence_length, dtype=jnp.int32), positions.shape) + row_indices = jnp.broadcast_to(jnp.arange(logits.shape[0], dtype=jnp.int32)[:, None], positions.shape) + safe_positions = jnp.clip(positions, 0, sequence_length - 1) + updates = jnp.where(validity_mask, array_indices + 1, 0) + position_to_index = jnp.zeros_like(positions).at[row_indices, safe_positions].max(updates) + previous_positions = jnp.maximum(positions - 1, 0) + source_indices = jnp.take_along_axis(position_to_index, jnp.clip(previous_positions, 0, sequence_length - 1), axis=1) + source_indices = jnp.maximum(source_indices - 1, 0) + return jnp.take_along_axis(logits, source_indices[..., None], axis=1) + raise ValueError(f"Unsupported block-diffusion logit alignment: {alignment}") diff --git a/src/maxtext/input_pipeline/hf_data_processing.py b/src/maxtext/input_pipeline/hf_data_processing.py index 370f1895bd..9b388778cd 100644 --- a/src/maxtext/input_pipeline/hf_data_processing.py +++ b/src/maxtext/input_pipeline/hf_data_processing.py @@ -42,6 +42,41 @@ def _get_pad_id(tokenizer): return pad_id +def _get_training_objective_transform( + config, + *, + shift, + use_dpo, + use_sft, + completion_only, + packing, + pad_id, + bos_token_id, +): + """Selects the aligned or next-token target transform for an HF pipeline.""" + objective = getattr(config, "training_objective", "causal_lm") + if objective == "block_diffusion": + if packing: + raise ValueError("Block-diffusion training requires packing=False.") + if use_dpo: + raise ValueError("Block-diffusion training is not compatible with DPO.") + return input_pipeline_utils.BlockDiffusionCorruption( + block_size=config.block_diffusion_block_size, + mask_id=config.block_diffusion_mask_id, + min_noise=config.block_diffusion_min_noise, + axis=1, + completion_only=bool(use_sft and completion_only), + seed_first_token=config.block_diffusion_canvas_policy == "seed_and_mask", + include_seed_in_loss=config.block_diffusion_logit_alignment == "shifted", + select_last_completion_suffix=getattr(config, "distill_data_source", "dataset") == "student_rollout", + ) + if objective != "causal_lm": + raise ValueError(f"Unsupported training objective: {objective}") + if shift and not use_dpo: + return input_pipeline_utils.ShiftData(ignored_ids=[pad_id, bos_token_id], axis=1) + return None + + def vision_sft_preprocessing_pipeline( dataset, config, @@ -321,12 +356,15 @@ def preprocessing_pipeline( ) operations = [] if use_sft: + is_block_diffusion = getattr(config, "training_objective", "causal_lm") == "block_diffusion" operations.append( input_pipeline_utils.SFTPromptMasking( text_column_name=data_column_names[0], completion_only=sft_train_on_completion_only, max_target_length=max_target_length, unk_id=pad_id, + target_aligned=is_block_diffusion, + emit_completion_mask=is_block_diffusion, ) ) data_column_names = ("inputs", "targets") @@ -357,8 +395,18 @@ 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: - operations.append(input_pipeline_utils.ShiftData(ignored_ids=[pad_id, tokenizer.bos_token_id], axis=1)) + target_transform = _get_training_objective_transform( + config, + shift=shift, + use_dpo=use_dpo, + use_sft=use_sft, + completion_only=sft_train_on_completion_only, + packing=packing, + pad_id=pad_id, + bos_token_id=tokenizer.bos_token_id, + ) + if target_transform is not None: + operations.append(target_transform) # Since HuggingFace IterableDataset does not support access through index # Indexes generated by dummy_index_sampler is not used. diff --git a/src/maxtext/input_pipeline/input_pipeline_utils.py b/src/maxtext/input_pipeline/input_pipeline_utils.py index d996a88908..9eead4ac91 100644 --- a/src/maxtext/input_pipeline/input_pipeline_utils.py +++ b/src/maxtext/input_pipeline/input_pipeline_utils.py @@ -325,14 +325,25 @@ def tokenization(example, hf_tokenizer, truncation, max_length, column_names): @dataclasses.dataclass class SFTPromptMasking(grain.MapTransform): """Construct inputs and targets for SFT training. Concat prompt and completion to generate inputs. - For targets, if train on completion only, the prompt will be masked by unk_id. Otherwise the same as inputs. + Causal SFT can mask prompt targets with ``unk_id``. Target-aligned objectives + keep clean targets and may emit role-derived completion metadata instead. """ - def __init__(self, text_column_name, completion_only, max_target_length, unk_id=0): + def __init__( + self, + text_column_name, + completion_only, + max_target_length, + unk_id=0, + target_aligned=False, + emit_completion_mask=False, + ): self.text_column_name = text_column_name self.completion_only = completion_only self.max_target_length = max_target_length self.unk_id = unk_id + self.target_aligned = target_aligned + self.emit_completion_mask = emit_completion_mask def map(self, element): """ @@ -340,18 +351,24 @@ def map(self, element): It concatenates the prompt and completion to form the `inputs` sequence. For the `targets` sequence: - If `self.completion_only` is `True`, the prompt portion of the - concatenated sequence is masked using `self.unk_id`. + concatenated sequence is masked using `self.unk_id`, unless + `self.target_aligned` requests clean same-position targets. - If `self.completion_only` is `False`, the target sequence is identical to the input sequence. """ - inputs, targets = [], [] + inputs, targets, completion_mask = [], [], [] for i, text in enumerate(element[self.text_column_name]): inputs += text - targets += [self.unk_id] * len(text) if self.completion_only and element["is_prompt"][i] else text - return { + is_prompt = element["is_prompt"][i] + targets += [self.unk_id] * len(text) if self.completion_only and is_prompt and not self.target_aligned else text + completion_mask += [not is_prompt] * len(text) + output = { "inputs": np.asarray(inputs[: self.max_target_length], dtype=np.int32), "targets": np.asarray(targets[: self.max_target_length], dtype=np.int32), } + if self.emit_completion_mask: + output["completion_mask"] = np.asarray(completion_mask[: self.max_target_length], dtype=np.int32) + return output @dataclasses.dataclass @@ -838,7 +855,7 @@ def map( """map to each element""" data_columns = list(element.keys()) for data_column in data_columns: - if data_column != "images": + if data_column not in ("images", "completion_mask"): if isinstance(element[data_column], mm_utils.PreprocessorOutput): raise TypeError("Only 'images' column can be of type PreprocessorOutput.") @@ -863,6 +880,8 @@ def map( element["images"] = self._pad_image_and_mask(element["images"]) # pyrefly: ignore[bad-argument-type] + elif key == "completion_mask": + element[key] = self._pad_text(element[key], self.max_length, 0) # pyrefly: ignore[bad-argument-type] elif "true_length" not in key: element[key] = self._pad_text(element[key], self.max_length, self.pad_id) # pyrefly: ignore[bad-argument-type] return element @@ -986,6 +1005,136 @@ def map(self, element): return shift_and_refine(element, ignored_ids=self.ignored_ids, axis=self.axis) +def _select_last_completion_suffix(validity_mask, completion_mask, axis): + """Selects the final assistant turn and truncates later conversation turns.""" + sequence_length = validity_mask.shape[axis] + moved_validity = np.moveaxis(validity_mask, axis, -1) + moved_completion = np.moveaxis(completion_mask, axis, -1) + flat_validity = moved_validity.reshape(-1, sequence_length) + flat_completion = moved_completion.reshape(-1, sequence_length) + positions = np.arange(sequence_length)[None, :] + + has_completion = flat_completion.any(axis=-1) + last_completion = np.max(np.where(flat_completion, positions, -1), axis=-1) + last_prefix_token = np.max( + np.where((positions < last_completion[:, None]) & ~flat_completion, positions, -1), + axis=-1, + ) + suffix_start = last_prefix_token + 1 + has_prompt = np.any(flat_validity & (positions < suffix_start[:, None]), axis=-1) + nonempty_rows = flat_validity.any(axis=-1) + invalid_rows = nonempty_rows & (~has_completion | ~has_prompt) + if np.any(invalid_rows): + raise ValueError("diffusion OPD requires a completion span after a prompt") + + retained = flat_validity & (positions <= last_completion[:, None]) & has_completion[:, None] + selected = retained & (positions >= suffix_start[:, None]) + selected = selected.reshape(moved_completion.shape) + retained = retained.reshape(moved_validity.shape) + return np.moveaxis(selected, -1, axis), np.moveaxis(retained, -1, axis) + + +@dataclasses.dataclass +class BlockDiffusionCorruption(grain.RandomMapTransform): + """Builds explicit completion, corruption, and loss masks for block diffusion.""" + + def __init__( + self, + block_size: int, + mask_id: int, + min_noise: float = 1.0e-3, + axis: int = 1, + completion_only: bool = False, + seed_first_token: bool = False, + include_seed_in_loss: bool = False, + select_last_completion_suffix: bool = False, + ): + if block_size <= 0: + raise ValueError(f"block_size must be positive, got {block_size}") + if not 0.0 < min_noise <= 1.0: + raise ValueError(f"min_noise must satisfy 0 < min_noise <= 1, got {min_noise}") + if include_seed_in_loss and not seed_first_token: + raise ValueError("include_seed_in_loss requires seed_first_token=True") + if select_last_completion_suffix and not completion_only: + raise ValueError("select_last_completion_suffix requires completion_only=True") + self.block_size = block_size + self.mask_id = mask_id + self.min_noise = min_noise + self.axis = axis + self.completion_only = completion_only + self.seed_first_token = seed_first_token + self.include_seed_in_loss = include_seed_in_loss + self.select_last_completion_suffix = select_last_completion_suffix + + def random_map(self, element, rng: np.random.Generator): + inputs = np.asarray(element["inputs"]) + targets = np.asarray(element["targets"]) + targets_segmentation = np.asarray(element["targets_segmentation"]) + if inputs.shape != targets.shape or inputs.shape != targets_segmentation.shape: + raise ValueError( + "inputs, targets, and targets_segmentation must have identical shapes, got " + f"{inputs.shape}, {targets.shape}, and {targets_segmentation.shape}" + ) + if inputs.ndim == 0 or not -inputs.ndim <= self.axis < inputs.ndim: + raise ValueError(f"axis {self.axis} is invalid for inputs with {inputs.ndim} dimensions") + + axis = self.axis % inputs.ndim + validity_mask = targets_segmentation != 0 + if self.completion_only and "completion_mask" not in element: + raise ValueError("completion_only block diffusion requires an explicit completion_mask") + completion_mask = np.asarray(element.get("completion_mask", validity_mask)) != 0 + if completion_mask.shape != inputs.shape: + raise ValueError(f"completion_mask must match inputs shape; received {completion_mask.shape} and {inputs.shape}") + completion_mask &= validity_mask + if self.select_last_completion_suffix: + completion_mask, validity_mask = _select_last_completion_suffix(validity_mask, completion_mask, axis) + inputs = np.where(validity_mask, inputs, inputs.dtype.type(0)) + targets = np.where(validity_mask, targets, targets.dtype.type(0)) + targets_segmentation = np.where( + validity_mask, targets_segmentation, targets_segmentation.dtype.type(0) + ) + eligible_mask = completion_mask if self.completion_only else validity_mask + + moved_eligible = np.moveaxis(eligible_mask, axis, -1) + rows = moved_eligible.reshape(-1, inputs.shape[axis]) + block_count = (rows.shape[-1] + self.block_size - 1) // self.block_size + padded_length = block_count * self.block_size + padded_rows = np.pad(rows, ((0, 0), (0, padded_length - rows.shape[-1]))) + eligible_blocks = padded_rows.reshape(rows.shape[0], block_count, self.block_size) + seed_loss_blocks = np.zeros_like(eligible_blocks) + if self.seed_first_token: + if self.include_seed_in_loss: + seed_loss_blocks = eligible_blocks & (np.arange(self.block_size) == 0) + seed_loss_blocks &= np.arange(block_count)[None, :, None] != 0 + eligible_blocks = eligible_blocks & (np.arange(self.block_size) != 0) + + noise = rng.uniform(self.min_noise, 1.0, size=eligible_blocks.shape[:-1] + (1,)) + selected_blocks = eligible_blocks & (rng.random(size=eligible_blocks.shape) < noise) + needs_fallback = eligible_blocks.any(axis=-1) & ~selected_blocks.any(axis=-1) + fallback_scores = np.where(eligible_blocks, rng.random(size=eligible_blocks.shape), -1.0) + fallback_offsets = np.argmax(fallback_scores, axis=-1) + fallback = np.arange(self.block_size) == fallback_offsets[..., None] + selected_blocks |= needs_fallback[..., None] & fallback + + corruption_rows = selected_blocks.reshape(rows.shape[0], padded_length)[:, : rows.shape[-1]] + loss_rows = (selected_blocks | seed_loss_blocks).reshape(rows.shape[0], padded_length)[:, : rows.shape[-1]] + corruption_mask = np.moveaxis(corruption_rows.reshape(moved_eligible.shape), -1, axis) + targets_loss_mask = np.moveaxis(loss_rows.reshape(moved_eligible.shape), -1, axis) + output = dict(element) + output["inputs"] = np.where(corruption_mask, inputs.dtype.type(self.mask_id), inputs) + output["targets"] = targets + if self.select_last_completion_suffix: + inputs_segmentation = np.asarray(element["inputs_segmentation"]) + output["inputs_segmentation"] = np.where( + validity_mask, inputs_segmentation, inputs_segmentation.dtype.type(0) + ) + output["targets_segmentation"] = targets_segmentation + output["completion_mask"] = completion_mask.astype(targets_segmentation.dtype) + output["corruption_mask"] = corruption_mask.astype(targets_segmentation.dtype) + output["targets_loss_mask"] = targets_loss_mask.astype(targets_segmentation.dtype) + return output + + @dataclasses.dataclass class ComputeQwen3OmniPositions(grain.MapTransform): """Computes 3D position IDs for Qwen3-Omni multimodal sequences. diff --git a/src/maxtext/input_pipeline/synthetic_data_processing.py b/src/maxtext/input_pipeline/synthetic_data_processing.py index d10816435c..a184e522cd 100644 --- a/src/maxtext/input_pipeline/synthetic_data_processing.py +++ b/src/maxtext/input_pipeline/synthetic_data_processing.py @@ -111,6 +111,11 @@ def get_place_holder_synthetic_data(config: pyconfig.HyperParameters): "targets_position": neg_ones, "targets_segmentation": neg_ones, } + if getattr(config, "training_objective", "causal_lm") == "block_diffusion": + zeros = np.zeros_like(neg_ones) + batch["completion_mask"] = zeros + batch["corruption_mask"] = zeros + batch["targets_loss_mask"] = zeros def infinite_iterator(): while True: diff --git a/src/maxtext/integration/tunix/diffusion_rl.py b/src/maxtext/integration/tunix/diffusion_rl.py new file mode 100644 index 0000000000..bfefcfe6bc --- /dev/null +++ b/src/maxtext/integration/tunix/diffusion_rl.py @@ -0,0 +1,459 @@ +# Copyright 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 +# +# http://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. + +"""Trajectory-aware MaxText adapters for block-diffusion reinforcement learning.""" + +from collections.abc import Sequence + +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.diffusion import scoring +from maxtext.diffusion import denoise +from tunix.diffusion import types as diffusion_types +from tunix.rl import reshard as rl_reshard +from tunix.rl.rollout import base_rollout + + +def _concrete_numpy(value): + if isinstance(value, jax.core.Tracer): + return None + if isinstance(value, jax.Array) and not value.is_fully_addressable: + return None + return np.asarray(value) + + +def prepare_diffusion_policy_batch( + *, + prompt_tokens, + prompt_mask, + completion_tokens, + completion_mask, + action_steps, + loss_mask=None, + inactive_target_id: int = 0, +) -> diffusion_types.DiffusionTokenBatch: + """Builds one completion-aligned policy batch from a sampled denoising trace. + + ``completion_mask`` describes the full canvas used during sampling, while + ``loss_mask`` selects actions owned by the RL objective. They intentionally + differ after an EOS token: exact replay still needs the post-EOS trace as + context for earlier bidirectional actions, but those later actions carry no + policy loss. + """ + expected_shape = tuple(completion_tokens.shape) + if tuple(completion_mask.shape) != expected_shape or tuple(action_steps.shape) != expected_shape: + raise ValueError("completion_tokens, completion_mask, and action_steps must have identical shapes") + if tuple(prompt_tokens.shape) != tuple(prompt_mask.shape): + raise ValueError("prompt_tokens and prompt_mask must have identical shapes") + if prompt_tokens.shape[0] != completion_tokens.shape[0]: + raise ValueError("prompt and completion batches must have the same leading dimension") + if loss_mask is None: + loss_mask = completion_mask + if tuple(loss_mask.shape) != expected_shape: + raise ValueError("loss_mask must match completion_tokens shape") + concrete_mask = _concrete_numpy(completion_mask) + concrete_steps = _concrete_numpy(action_steps) + concrete_loss_mask = _concrete_numpy(loss_mask) + if concrete_mask is not None and concrete_steps is not None and concrete_loss_mask is not None: + active = np.asarray(concrete_mask, dtype=bool) + weighted = np.asarray(concrete_loss_mask, dtype=bool) + concrete_steps = np.asarray(concrete_steps) + if not np.array_equal(concrete_steps >= 0, active): + raise ValueError("each active completion token must have exactly one nonnegative action step") + if np.any(weighted & ~active): + raise ValueError("loss_mask must be a subset of the sampled completion canvas") + if np.any(concrete_steps[active] >= completion_tokens.shape[1]): + raise ValueError("active action steps must be smaller than the completion length") + use_numpy = all( + isinstance(value, np.ndarray) + for value in (prompt_tokens, prompt_mask, completion_tokens, completion_mask, action_steps, loss_mask) + ) + array_module = np if use_numpy else jnp + completion_tokens = array_module.asarray(completion_tokens) + loss_mask = array_module.asarray(loss_mask, dtype=array_module.bool_) + target_ids = array_module.where( + loss_mask, + completion_tokens, + array_module.asarray(inactive_target_id, dtype=completion_tokens.dtype), + ) + return diffusion_types.DiffusionTokenBatch.create( + model_inputs={ + "prompt_tokens": array_module.asarray(prompt_tokens), + "prompt_mask": array_module.asarray(prompt_mask, dtype=array_module.bool_), + "completion_tokens": completion_tokens, + "completion_mask": array_module.asarray(completion_mask, dtype=array_module.bool_), + "action_steps": array_module.asarray(action_steps, dtype=array_module.int32), + }, + target_ids=target_ids, + loss_weights=array_module.asarray(loss_mask, dtype=array_module.float32), + ) + + +def _logical_positions(validity_mask): + positions = jnp.cumsum(jnp.asarray(validity_mask, dtype=jnp.int32), axis=1) - 1 + return jnp.where(validity_mask, positions, 0) + + +def make_diffusion_trace_logits_fn(config): + """Builds a scorer that replays the exact pre-commit canvas for each action.""" + mask_id = int(config.block_diffusion_mask_id) + alignment = config.block_diffusion_logit_alignment + vocab_size = int(config.vocab_size) + + def logits_fn(model: nnx.Module, model_inputs: diffusion_types.ModelInputs): + prompt_tokens = jnp.asarray(model_inputs["prompt_tokens"]) + prompt_mask = jnp.asarray(model_inputs["prompt_mask"], dtype=jnp.bool_) + completion_tokens = jnp.asarray(model_inputs["completion_tokens"]) + completion_mask = jnp.asarray(model_inputs["completion_mask"], dtype=jnp.bool_) + action_steps = jnp.asarray(model_inputs["action_steps"], dtype=jnp.int32) + batch_size, completion_length = completion_tokens.shape + prompt_length = prompt_tokens.shape[1] + validity_mask = jnp.concatenate([prompt_mask, completion_mask], axis=1) + positions = _logical_positions(validity_mask) + segment_ids = validity_mask.astype(jnp.int32) + base_model = getattr(model, "base", model) + graphdef, params, rest = nnx.split(base_model, nnx.Param, ...) + output_logits = jnp.zeros((batch_size, completion_length, vocab_size), dtype=jnp.float32) + + def score_action_step(step, accumulated_logits): + active_actions = completion_mask & (action_steps == step) + + def run_forward(current_logits): + pre_commit_completion = jnp.where(action_steps >= step, mask_id, completion_tokens) + canvas = jnp.concatenate([prompt_tokens, pre_commit_completion], axis=1) + local_model = nnx.merge(graphdef, params, rest, copy=True) + raw_logits = local_model( + decoder_input_tokens=canvas, + decoder_positions=positions, + decoder_segment_ids=segment_ids, + enable_dropout=False, + decoder_target_tokens=canvas, + decoder_target_mask=segment_ids, + ) + target_aligned = scoring.align_logits_to_targets(raw_logits, alignment, positions, validity_mask) + completion_logits = jnp.asarray(target_aligned[:, prompt_length:, :], dtype=jnp.float32) + completion_logits = completion_logits.at[..., mask_id].set(-jnp.inf) + return jnp.where(active_actions[..., None], completion_logits, current_logits) + + return jax.lax.cond(jnp.any(active_actions), run_forward, lambda value: value, accumulated_logits) + + output_logits = jax.lax.fori_loop(0, completion_length, score_action_step, output_logits) + return jnp.where(completion_mask[..., None], output_logits, 0.0) + + return logits_fn + + +def make_diffusion_rollout_fn(config, temperature: float, *, sample_tokens: bool = True): + """Builds one compiled stochastic rollout over the current MaxText policy.""" + max_denoise_steps = config.rl.diffusion_max_denoise_steps + if max_denoise_steps == -1: + max_denoise_steps = config.block_diffusion_block_size + + @nnx.jit + def rollout(model, initial_tokens, positions, segment_ids, completion_mask, rng): + graphdef, params, rest = nnx.split(model, nnx.Param, ...) + validity_mask = segment_ids != 0 + + def canvas_logits(canvas): + local_model = nnx.merge(graphdef, params, rest, copy=True) + base_model = getattr(local_model, "base", local_model) + raw_logits = base_model( # pylint: disable=not-callable + decoder_input_tokens=canvas, + decoder_positions=positions, + decoder_segment_ids=segment_ids, + enable_dropout=False, + decoder_target_tokens=canvas, + decoder_target_mask=segment_ids, + ) + return scoring.align_logits_to_targets( + raw_logits, + config.block_diffusion_logit_alignment, + positions, + validity_mask, + ) + + return denoise.low_confidence_rollout( + canvas_logits, + initial_tokens, + positions, + validity_mask, + completion_mask, + block_size=config.block_diffusion_block_size, + mask_id=config.block_diffusion_mask_id, + logit_alignment=config.block_diffusion_logit_alignment, + canvas_policy=config.block_diffusion_canvas_policy, + confidence_threshold=config.rl.diffusion_confidence_threshold, + temperature=temperature, + max_denoise_steps=max_denoise_steps, + rng=rng if sample_tokens else None, + ) + + return rollout + + +def resolve_stop_token_ids(config, tokenizer_eos_id: int | None) -> tuple[int, ...]: + """Resolves and validates the model's generated completion terminators.""" + configured = tuple(int(token_id) for token_id in config.rl.diffusion_stop_token_ids) + if configured: + stop_token_ids = configured + elif tokenizer_eos_id is not None: + stop_token_ids = (int(tokenizer_eos_id),) + else: + raise ValueError("block-diffusion RL requires an EOS token or explicit diffusion_stop_token_ids") + if len(set(stop_token_ids)) != len(stop_token_ids) or any( + token_id < 0 or token_id >= config.vocab_size for token_id in stop_token_ids + ): + raise ValueError("diffusion stop-token IDs must be unique and inside the model vocabulary") + if config.block_diffusion_mask_id in stop_token_ids: + raise ValueError("the diffusion mask token cannot also be a generated stop token") + return stop_token_ids + + +def policy_mask_at_stop( + tokens: jax.Array, + action_steps: jax.Array, + stop_token_ids: Sequence[int], +) -> jax.Array: + """Selects sampled actions through the first generated stop token.""" + active = action_steps >= 0 + stop_positions = active & jnp.isin(tokens, jnp.asarray(stop_token_ids, tokens.dtype)) + seen_stop = jax.lax.associative_scan(jnp.logical_or, stop_positions, axis=1) + after_stop = jnp.concatenate([jnp.zeros_like(seen_stop[:, :1]), seen_stop[:, :-1]], axis=1) + return active & ~after_stop + + +class MaxTextDiffusionRollout(base_rollout.BaseRollout): + """Correctness-first in-process rollout that exports exact denoising traces.""" + + def __init__( + self, + rollout_actor, + tokenizer, + mesh, + rollout_config, + maxtext_config, + cache_config_or_size=None, + ): + del cache_config_or_size + self._model = rollout_actor + self._tokenizer = tokenizer + self._mesh = mesh + self._rollout_config = rollout_config + self._config = maxtext_config + self._pad_id = tokenizer.pad_id() + if self._pad_id is None: + raise ValueError("block-diffusion RL requires a tokenizer pad token") + self._pad_id = int(self._pad_id) + if not 0 <= self._pad_id < maxtext_config.vocab_size: + raise ValueError("the tokenizer pad token must be inside the model vocabulary") + self._eos_id = tokenizer.eos_id() + self._stop_token_ids = resolve_stop_token_ids(maxtext_config, self._eos_id) + self._rollout_fns = { + (float(rollout_config.temperature), True): make_diffusion_rollout_fn( + maxtext_config, float(rollout_config.temperature) + ) + } + configured_seed = getattr(rollout_config, "seed", None) + if configured_seed is None: + self._base_rng = jax.random.PRNGKey(int(getattr(maxtext_config, "data_shuffle_seed", 0))) + else: + configured_seed = jnp.asarray(configured_seed) + if not configured_seed.shape: + self._base_rng = jax.random.PRNGKey(int(configured_seed)) + elif tuple(configured_seed.shape) == (2,): + self._base_rng = configured_seed + else: + raise ValueError("RolloutConfig.seed must be an integer or one JAX PRNG key") + training_stream = self._sampling_stream_key(rollout_config) + self._rng_streams = {training_stream: self._base_rng} + self._generation_context = None + self._generation_context_invocation = 0 + self._generation_call = 0 + + @property + def mesh(self): + return self._mesh + + def _encode_left_padded(self, prompts, max_prompt_length): + rows = [] + masks = [] + for prompt in prompts: + token_ids = list(self._tokenizer.encode(prompt))[-max_prompt_length:] + padding = max_prompt_length - len(token_ids) + rows.append([self._pad_id] * padding + token_ids) + masks.append([False] * padding + [True] * len(token_ids)) + return np.asarray(rows, dtype=np.int32), np.asarray(masks, dtype=bool) + + def _rollout_fn_for(self, rollout_config): + """Returns a cached rollout compiled for the requested sampling mode.""" + top_k = getattr(rollout_config, "top_k", None) + top_p = getattr(rollout_config, "top_p", 1.0) + if top_p not in (None, 1.0) or top_k not in (None, -1, 1): + raise ValueError("block-diffusion rollout supports categorical sampling or greedy top_k=1") + sample_tokens = top_k != 1 + key = (float(rollout_config.temperature), sample_tokens) + if key not in self._rollout_fns: + self._rollout_fns[key] = make_diffusion_rollout_fn( + self._config, + float(rollout_config.temperature), + sample_tokens=sample_tokens, + ) + return self._rollout_fns[key] + + def _sampling_stream_key(self, rollout_config): + return ( + float(rollout_config.temperature), + getattr(rollout_config, "top_k", None), + getattr(rollout_config, "top_p", 1.0), + bool(getattr(rollout_config, "return_logprobs", False)), + ) + + def _next_rng(self, rollout_config): + """Advances the deterministic RNG stream for one generation call.""" + if self._generation_context is not None: + global_step, mode = self._generation_context + mode_id = 0 if mode == "train" else 1 + rollout_key = jax.random.fold_in(self._base_rng, global_step) + rollout_key = jax.random.fold_in(rollout_key, mode_id) + rollout_key = jax.random.fold_in(rollout_key, self._generation_context_invocation) + rollout_key = jax.random.fold_in(rollout_key, self._generation_call) + self._generation_call += 1 + return rollout_key + stream = self._sampling_stream_key(rollout_config) + if stream not in self._rng_streams: + self._rng_streams[stream] = jax.random.fold_in(self._base_rng, len(self._rng_streams)) + self._rng_streams[stream], rollout_key = jax.random.split(self._rng_streams[stream]) + return rollout_key + + def set_generation_context(self, *, global_step: int, mode): + """Derives resumable rollout keys without coupling train and eval calls.""" + context = (int(global_step), str(mode)) + if context == self._generation_context: + self._generation_context_invocation += 1 + else: + self._generation_context = context + self._generation_context_invocation = 0 + self._generation_call = 0 + + def generate(self, prompts, rollout_config, **kwargs): + del kwargs + rollout_fn = self._rollout_fn_for(rollout_config) + prompt_tokens, prompt_mask = self._encode_left_padded(prompts, rollout_config.max_prompt_length) + batch_size = prompt_tokens.shape[0] + completion_length = rollout_config.max_tokens_to_generate + completion_mask = np.ones((batch_size, completion_length), dtype=bool) + validity_mask = np.concatenate([prompt_mask, completion_mask], axis=1) + positions = _logical_positions(jnp.asarray(validity_mask)) + segment_ids = jnp.asarray(validity_mask, dtype=jnp.int32) + initial_tokens = jnp.concatenate( + [ + jnp.asarray(prompt_tokens), + jnp.full( + (batch_size, completion_length), + self._config.block_diffusion_mask_id, + dtype=jnp.int32, + ), + ], + axis=1, + ) + full_completion_mask = jnp.concatenate( + [ + jnp.zeros(prompt_tokens.shape, dtype=jnp.bool_), + jnp.ones((batch_size, completion_length), dtype=jnp.bool_), + ], + axis=1, + ) + rollout_key = self._next_rng(rollout_config) + trace = rollout_fn( + self._model, + initial_tokens, + positions, + segment_ids, + full_completion_mask, + rollout_key, + ) + completion_tokens = trace.tokens[:, -completion_length:] + completion_steps = trace.action_steps[:, -completion_length:] + completion_logps = trace.action_logps[:, -completion_length:] + output_mask = policy_mask_at_stop( + completion_tokens, + completion_steps, + self._stop_token_ids, + ) + for name, value in ( + ("completion tokens", completion_tokens), + ("action steps", completion_steps), + ("action log probabilities", completion_logps), + ("output mask", output_mask), + ): + if isinstance(value, jax.Array) and not value.is_fully_addressable: + raise RuntimeError( + f"block-diffusion RL requires fully addressable {name}; use the supported Pathways single-controller mode" + ) + host_completion_tokens = np.asarray(jax.device_get(completion_tokens)) + host_completion_steps = np.asarray(jax.device_get(completion_steps)) + host_completion_logps = np.asarray(jax.device_get(completion_logps)) + host_completion_mask = host_completion_steps >= 0 + host_output_mask = np.asarray(jax.device_get(output_mask)) + if not np.all(host_completion_mask) or np.any(host_completion_tokens == self._config.block_diffusion_mask_id): + raise RuntimeError("block-diffusion rollout ended with unresolved mask tokens") + if not np.all(np.isfinite(host_completion_logps[host_completion_mask])): + raise RuntimeError("block-diffusion rollout produced non-finite action log probabilities") + diffusion_batch = prepare_diffusion_policy_batch( + prompt_tokens=prompt_tokens, + prompt_mask=prompt_mask, + completion_tokens=host_completion_tokens, + completion_mask=host_completion_mask, + action_steps=host_completion_steps, + inactive_target_id=self._pad_id, + ) + tokens_out = [] + text_out = [] + logps_out = [] + for tokens_row, active_row, logps_row in zip(host_completion_tokens, host_output_mask, host_completion_logps): + length = int(active_row.sum()) + trimmed_tokens = np.asarray(tokens_row[:length], dtype=np.int32) + tokens_out.append(trimmed_tokens) + text_out.append(self._tokenizer.decode(trimmed_tokens.tolist())) + logps_out.append(np.asarray(logps_row, dtype=np.float32)) + return base_rollout.RolloutOutput( + text=text_out, + logits=None, + tokens=tokens_out, + left_padded_prompt_tokens=prompt_tokens, + logprobs=logps_out if rollout_config.return_logprobs else None, + diffusion_batch=diffusion_batch, + ) + + def update_params(self, params, filter_types=None): + destination = nnx.state(self._model, filter_types) if filter_types is not None else nnx.state(self._model) + nnx.update(self._model, rl_reshard.reshard_pytree(params, destination)) + + def get_per_token_logps(self, prompt_tokens, completion_tokens): + del prompt_tokens, completion_tokens + raise NotImplementedError( + "block-diffusion policy scores require RolloutOutput.diffusion_batch; final-sequence AR scoring is invalid" + ) + + def pad_id(self): + return self._pad_id + + def eos_id(self): + return self._eos_id + + def model(self): + return self._model diff --git a/src/maxtext/integration/tunix/diffusion_sft.py b/src/maxtext/integration/tunix/diffusion_sft.py new file mode 100644 index 0000000000..b85132db38 --- /dev/null +++ b/src/maxtext/integration/tunix/diffusion_sft.py @@ -0,0 +1,180 @@ +# Copyright 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 +# +# http://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. + +"""MaxText adapters for target-aligned Tunix diffusion SFT.""" + +from collections.abc import Mapping +from typing import Any + +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.diffusion import scoring +from tunix.diffusion import types as diffusion_types + + +_REQUIRED_FIELDS = ( + "inputs", + "inputs_position", + "inputs_segmentation", + "targets", + "targets_position", + "targets_segmentation", + "completion_mask", + "corruption_mask", + "targets_loss_mask", +) + + +def _concrete_numpy(value): + if isinstance(value, jax.core.Tracer): + return None + if isinstance(value, jax.Array) and not value.is_fully_addressable: + return None + return np.asarray(value) + + +def _validate_batch_masks( + positions, + validity_mask, + completion_mask, + corruption_mask, + loss_weights, + *, + alignment, + block_size, +): + shapes = { + "positions": tuple(positions.shape), + "validity_mask": tuple(validity_mask.shape), + "completion_mask": tuple(completion_mask.shape), + "corruption_mask": tuple(corruption_mask.shape), + "targets_loss_mask": tuple(loss_weights.shape), + } + if len(set(shapes.values())) != 1: + raise ValueError( + f"diffusion SFT masks must have identical shapes; received {shapes}" + ) + concrete = [ + _concrete_numpy(value) + for value in ( + positions, + validity_mask, + completion_mask, + corruption_mask, + loss_weights, + ) + ] + if any(value is None for value in concrete): + return + concrete_positions = np.asarray(concrete[0]) + validity, completion, corruption = ( + np.asarray(value, dtype=bool) for value in concrete[1:4] + ) + weights = np.asarray(concrete[4]) + weighted = weights != 0 + if not np.all(np.isfinite(weights)) or np.any(weights < 0): + raise ValueError("targets_loss_mask must contain finite nonnegative weights") + if np.any(completion & ~validity): + raise ValueError("completion_mask must be a subset of valid target positions") + if np.any(corruption & ~completion): + raise ValueError("corruption_mask must be a subset of completion_mask") + if np.any(weighted & ~completion): + raise ValueError("targets_loss_mask must be a subset of completion_mask") + allowed = corruption.copy() + if alignment == "shifted": + allowed |= (concrete_positions > 0) & (concrete_positions % block_size == 0) + if np.any(weighted & ~allowed): + raise ValueError( + "diffusion SFT loss weights must own corrupted targets or shifted block anchors" + ) + + +def create_batch_adapter(config): + """Builds a raw MaxText batch adapter for Tunix diffusion SFT.""" + alignment = config.block_diffusion_logit_alignment + block_size = int(config.block_diffusion_block_size) + + def adapt(raw_batch: Mapping[str, Any]) -> diffusion_types.DiffusionTokenBatch: + missing = sorted(name for name in _REQUIRED_FIELDS if name not in raw_batch) + if missing: + raise ValueError( + f"block-diffusion SFT requires explicit batch fields; missing {missing}" + ) + targets = jnp.asarray(raw_batch["targets"]) + positions = jnp.asarray(raw_batch["targets_position"], dtype=jnp.int32) + validity_mask = jnp.asarray(raw_batch["targets_segmentation"]) != 0 + completion_mask = jnp.asarray(raw_batch["completion_mask"], dtype=jnp.bool_) + corruption_mask = jnp.asarray(raw_batch["corruption_mask"], dtype=jnp.bool_) + raw_loss_weights = jnp.asarray( + raw_batch["targets_loss_mask"], dtype=jnp.float32 + ) + _validate_batch_masks( + positions, + validity_mask, + completion_mask, + corruption_mask, + raw_loss_weights, + alignment=alignment, + block_size=block_size, + ) + allowed = corruption_mask + if alignment == "shifted": + allowed |= (positions > 0) & (positions % block_size == 0) + loss_weights = jnp.where( + validity_mask & completion_mask & allowed, raw_loss_weights, 0.0 + ) + return diffusion_types.DiffusionTokenBatch.create( + model_inputs={ + "input_tokens": jnp.asarray(raw_batch["inputs"]), + "input_positions": jnp.asarray( + raw_batch["inputs_position"], dtype=jnp.int32 + ), + "input_segmentation": jnp.asarray(raw_batch["inputs_segmentation"]), + "targets": targets, + "target_positions": positions, + "target_segmentation": jnp.asarray(raw_batch["targets_segmentation"]), + }, + target_ids=targets, + loss_weights=loss_weights, + ) + + return adapt + + +def create_target_aligned_logits_fn(config): + """Builds a MaxText scorer satisfying Tunix's diffusion logits contract.""" + alignment = config.block_diffusion_logit_alignment + enable_dropout = bool(config.enable_dropout) + + def logits_fn(model: nnx.Module, model_inputs: diffusion_types.ModelInputs): + base_model = getattr(model, "base", model) + logits = base_model( + decoder_input_tokens=model_inputs["input_tokens"], + decoder_positions=model_inputs["input_positions"], + decoder_segment_ids=model_inputs["input_segmentation"], + enable_dropout=enable_dropout, + decoder_target_tokens=model_inputs["targets"], + decoder_target_mask=model_inputs["target_segmentation"], + ) + return scoring.align_logits_to_targets( + logits, + alignment, + model_inputs["target_positions"], + model_inputs["target_segmentation"] != 0, + ) + + return logits_fn diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 8e6c96a72e..af48802cc9 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -84,6 +84,16 @@ dynamic_vector_slice_in_dim = jax.vmap(lax.dynamic_slice_in_dim, in_axes=(None, 0, None, None)) +def _resolve_attention_type(config: Config, attention_type: AttentionType | str | None) -> AttentionType: + configured_attention_type = AttentionType(getattr(config, "attention_type", AttentionType.GLOBAL.value)) + if attention_type is None: + return configured_attention_type + resolved_attention_type = AttentionType(attention_type) + if configured_attention_type == AttentionType.BLOCK_DIFFUSION and resolved_attention_type == AttentionType.GLOBAL: + return configured_attention_type + return resolved_attention_type + + def validate_compute_axis_order(s: AxisIdxes) -> None: valid_compute_axis_order = ((0, 1, 2, 3), (0, 2, 1, 3)) if s not in valid_compute_axis_order: # currently supported compute_axis_order @@ -204,6 +214,50 @@ def __hash__(self): ) +class BlockCausalMask(splash_attention_mask._ComputableMask): # pylint: disable=protected-access,abstract-method + """Lazy mask with bidirectional attention within causal blocks.""" + + block_size: int + + def __init__( + self, + shape: tuple[int, int], + block_size: int, + shard_count: int = 1, + ): + if block_size <= 0: + raise ValueError("block_size must be positive") + self.block_size = block_size + + def block_causal_mask_function(q_ids, kv_ids): + return (q_ids // self.block_size) >= (kv_ids // self.block_size) + + super().__init__( + shape=shape, + mask_function=block_causal_mask_function, + shard_count=shard_count, + ) + + def __eq__(self, other: object): + if not isinstance(other, type(self)): + return NotImplemented + return ( + self.shape == other.shape + and self.block_size == other.block_size + and np.array_equal(self.q_sequence, other.q_sequence) + ) + + def __hash__(self): + return hash( + ( + type(self), + self.shape, + self.block_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. @@ -234,6 +288,16 @@ def _generate_chunk_attention_mask(mask_shape: tuple[int, int], chunk_size: int, return chunk_mask +def _generate_block_causal_attention_mask(mask_shape: tuple[int, int], block_size: int, q_offset: int = 0) -> jax.Array: + """Generates a block-causal mask, bidirectional within each block.""" + if block_size <= 0: + raise ValueError("block_size must be positive") + + row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0) + q_offset + col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1) + return (row_ids // block_size) >= (col_ids // block_size) + + def _make_block_mask_indices(bidirectional_mask): """Creates block mask identifying segments based on a bidirectional mask. @@ -479,7 +543,13 @@ def __init__( self.dtype = dtype self.quant = quant self.kv_quant = kv_quant - self.attention_type = attention_type + self.attention_type = _resolve_attention_type(self.config, attention_type) + self.block_diffusion_block_size = getattr(self.config, "block_diffusion_block_size", None) + if self.attention_type == AttentionType.BLOCK_DIFFUSION: + if self.block_diffusion_block_size is None or self.block_diffusion_block_size <= 0: + raise ValueError("block_diffusion_block_size must be positive for block-diffusion attention") + if self.attention_kernel not in ("autoselected", "dot_product", "flash"): + raise ValueError("Block-diffusion attention is supported only by dot_product attention and TPU Splash attention.") # Block sizes are only used by TPU splash attention kernels. Exclude non-splash kernels if self.attention_kernel not in ( "dot_product", @@ -782,6 +852,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, :, :] @@ -861,6 +932,19 @@ def generate_attention_mask( ) output_mask = chunk_mask * output_mask + elif self.attention_type == AttentionType.BLOCK_DIFFUSION and model_mode != MODEL_MODE_AUTOREGRESSIVE: + if use_segment_positions: + block_mask = ( + (position_row_ids // self.block_diffusion_block_size) >= (position_col_ids // self.block_diffusion_block_size) + )[:, None, None, :, :] + else: + block_mask = _generate_block_causal_attention_mask( + mask_shape=(q_seq_len, kv_seq_len), + block_size=self.block_diffusion_block_size, + q_offset=next_pos, + )[None, None, None, :, :] + output_mask = block_mask if output_mask is None else jnp.logical_and(output_mask, block_mask) + if bidirectional_mask is not None: image_mask = _make_bidirectional_block_mask(bidirectional_mask) output_mask = output_mask | image_mask[:, None, None, ...] @@ -1126,6 +1210,11 @@ def apply_attention( return out, None, None else: + if self.attention_type == AttentionType.BLOCK_DIFFUSION: + raise ValueError( + "Block-diffusion flash attention is supported only by TPU Splash; " + "use attention='dot_product' on other hardware." + ) if model_mode == MODEL_MODE_AUTOREGRESSIVE: # fallback to dot_product as pallas gpu flash attention doesn't support decode stage return self.apply_attention_dot( @@ -1430,13 +1519,24 @@ def create_sa_config(config, query, key, attn_logits_soft_cap): sa_config = create_sa_config(self.config, query, key, attn_logits_soft_cap) mask_shape = (query.shape[2], key.shape[2]) # (q_seq_len, kv_seq_len) mask_module = tokamax_splash_mask if self.config.use_tokamax_splash else splash_attention_mask + use_load_balanced_cp = cp_size > 1 and load_balanced_context_parallel if self.attention_type == AttentionType.FULL: mask = mask_module.FullMask(mask_shape) + elif self.attention_type == AttentionType.BLOCK_DIFFUSION: + mask_type = LoadBalancedBlockCausalMask if use_load_balanced_cp else BlockCausalMask + mask_kwargs = {"cp_size": cp_size} if use_load_balanced_cp else {} + mask = mask_type( + shape=mask_shape, + block_size=self.block_diffusion_block_size, + **mask_kwargs, + ) 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. @@ -2447,3 +2547,21 @@ def __init__( shard_count=shard_count, ) self.q_sequence = _load_balanced_q_sequence(shape, cp_size) + + +class LoadBalancedBlockCausalMask(BlockCausalMask): # pylint: disable=abstract-method + """Lazy block-causal mask with load-balanced query positions.""" + + def __init__( + self, + shape: tuple[int, int], + block_size: int, + cp_size: int, + shard_count: int = 1, + ): + super().__init__( + shape=shape, + block_size=block_size, + shard_count=shard_count, + ) + self.q_sequence = _load_balanced_q_sequence(shape, cp_size) diff --git a/src/maxtext/layers/attentions.py b/src/maxtext/layers/attentions.py index 5a65b35ba4..32f233cbeb 100644 --- a/src/maxtext/layers/attentions.py +++ b/src/maxtext/layers/attentions.py @@ -51,7 +51,7 @@ AttentionType, ) from maxtext.layers import nnx_wrappers -from maxtext.layers.attention_op import AttentionOp +from maxtext.layers.attention_op import AttentionOp, _resolve_attention_type from maxtext.layers.embeddings import ( LLaMARotaryEmbedding, LlamaVisionRotaryEmbedding, @@ -120,7 +120,7 @@ def attention_as_linen( float32_logits: bool = False, # cast logits in float32 for stability. quant: Optional[Quant] = None, kv_quant: Optional[KVQuant] = None, - attention_type: AttentionType = AttentionType.GLOBAL, # Default to global attention + attention_type: AttentionType | None = None, attn_logits_soft_cap: float | None = None, sliding_window_size: int | None = None, use_ragged_attention: bool = False, @@ -277,7 +277,7 @@ def __init__( float32_logits: bool = False, # cast logits in float32 for stability. quant: Optional[Quant] = None, kv_quant: Optional[KVQuant] = None, - attention_type: AttentionType = AttentionType.GLOBAL, # Default to global attention + attention_type: AttentionType | None = None, attn_logits_soft_cap: float | None = None, sliding_window_size: int | None = None, use_ragged_attention: bool = False, @@ -388,7 +388,7 @@ def __init__( self.float32_logits = float32_logits self.quant = quant self.kv_quant = kv_quant - self.attention_type = attention_type + self.attention_type = _resolve_attention_type(self.config, attention_type) self.attn_logits_soft_cap = attn_logits_soft_cap self.sliding_window_size = sliding_window_size self.use_ragged_attention = use_ragged_attention diff --git a/src/maxtext/trainers/post_train/distillation/diffusion_opd.py b/src/maxtext/trainers/post_train/distillation/diffusion_opd.py new file mode 100644 index 0000000000..1a858ce9b5 --- /dev/null +++ b/src/maxtext/trainers/post_train/distillation/diffusion_opd.py @@ -0,0 +1,622 @@ +# Copyright 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 +# +# http://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. + +"""MaxText model-aware preparation for Tunix diffusion OPD.""" + +import math + +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.diffusion import denoise +from maxtext.diffusion import scoring +from maxtext.trainers.post_train.distillation import distillation_utils +from maxtext.utils import max_logging +from tunix.diffusion import types as diffusion_types +from tunix.distillation import diffusion as diffusion_distillation +from tunix.distillation import diffusion_opd as tunix_diffusion_opd +from tunix.sft import peft_trainer + + +def _resolve_logit_alignment(config) -> str: + if getattr(config, "training_objective", "causal_lm") == "block_diffusion": + return config.block_diffusion_logit_alignment + return "shifted" + + +def create_target_aligned_logits_fn(config, *, enable_dropout: bool): + """Builds a MaxText forward adapter satisfying Tunix's logits contract.""" + alignment = _resolve_logit_alignment(config) + + def logits_fn(model, model_inputs): + logits = model( + decoder_input_tokens=model_inputs["input_tokens"], + decoder_positions=model_inputs["positions"], + decoder_segment_ids=model_inputs["decoder_segment_ids"], + enable_dropout=enable_dropout, + decoder_target_tokens=model_inputs["targets"], + decoder_target_mask=model_inputs["targets_segmentation"], + ) + return scoring.align_logits_to_targets( + logits, + alignment, + model_inputs["positions"], + model_inputs["targets_segmentation"] != 0, + ) + + return logits_fn + + +def create_rollout_fn(config): + """Builds one compiled student rollout with deterministic dropout settings.""" + logits_fn = create_target_aligned_logits_fn(config, enable_dropout=False) + max_denoise_steps = config.distill_rollout_max_denoise_steps + if max_denoise_steps == -1: + max_denoise_steps = config.block_diffusion_block_size + + @nnx.jit + def rollout(model, initial_tokens, positions, decoder_segment_ids, completion_mask): + graphdef, params, rest = nnx.split(model, nnx.Param, ...) + validity_mask = decoder_segment_ids != 0 + + def canvas_logits(canvas): + # The denoising state machine uses nested JAX control flow. Reconstructing + # a local model at that trace level avoids mutating outer NNX Variables; + # rollout uses dropout=False, so its transient non-parameter state is + # intentionally discarded after each proposal forward. + local_model = nnx.merge(graphdef, params, rest, copy=True) + return logits_fn( + local_model, + { + "input_tokens": canvas, + "positions": positions, + "decoder_segment_ids": decoder_segment_ids, + "targets": canvas, + "targets_segmentation": decoder_segment_ids, + }, + ) + + return denoise.low_confidence_generate( + canvas_logits, + initial_tokens, + positions, + validity_mask, + completion_mask, + block_size=config.block_diffusion_block_size, + mask_id=config.block_diffusion_mask_id, + logit_alignment=config.block_diffusion_logit_alignment, + canvas_policy=config.block_diffusion_canvas_policy, + confidence_threshold=config.distill_rollout_confidence_threshold, + temperature=config.distill_rollout_temperature, + max_denoise_steps=max_denoise_steps, + ) + + return rollout + + +def create_teacher_score_fn(config): + """Builds a compiled frozen-teacher scoring function.""" + logits_fn = create_target_aligned_logits_fn(config, enable_dropout=False) + + @nnx.jit + def score(model, model_inputs): + graphdef, params, rest = nnx.split(model, nnx.Param, ...) + local_model = nnx.merge(graphdef, params, rest, copy=True) + return logits_fn(local_model, model_inputs) + + return score + + +def _concrete_numpy(value): + if isinstance(value, jax.core.Tracer): + return None + if isinstance(value, jax.Array) and not value.is_fully_addressable: + return None + return np.asarray(value) + + +def _validate_input_masks( + positions, + validity_mask, + completion_mask, + corruption_mask, + loss_mask, + *, + logit_alignment, + block_size, +): + """Validates role and loss ownership before launching model computation.""" + shapes = { + "positions": tuple(positions.shape), + "validity_mask": tuple(validity_mask.shape), + "completion_mask": tuple(completion_mask.shape), + "corruption_mask": tuple(corruption_mask.shape), + "targets_loss_mask": tuple(loss_mask.shape), + } + if len(set(shapes.values())) != 1: + raise ValueError(f"diffusion OPD masks must have identical shapes; received {shapes}") + concrete = [_concrete_numpy(value) for value in (positions, validity_mask, completion_mask, corruption_mask, loss_mask)] + if any(value is None for value in concrete): + return + concrete_positions = np.asarray(concrete[0]) + validity, completion, corruption, weights = (np.asarray(value, dtype=bool) for value in concrete[1:]) + if np.any(completion & ~validity): + raise ValueError("completion_mask must be a subset of valid target positions") + if np.any(corruption & ~completion): + raise ValueError("corruption_mask must be a subset of completion_mask") + if np.any(weights & ~completion): + raise ValueError("targets_loss_mask must be a subset of completion_mask") + if np.any(corruption & ~weights): + raise ValueError("every corrupted position must carry a positive target loss weight") + safe_clean_targets = np.zeros_like(weights) + if logit_alignment == "shifted": + safe_clean_targets = (concrete_positions > 0) & (concrete_positions % block_size == 0) + if np.any(weights & ~corruption & ~safe_clean_targets): + raise ValueError( + "weighted clean student targets are allowed only at shifted logical block anchors; " + "same-position targets must be corrupted" + ) + if np.any(weights & (concrete_positions == 0)): + raise ValueError("causal teacher distillation must not weight logical position zero") + + +def _safe_loss_weights( + positions, validity_mask, completion_mask, corruption_mask, loss_weights, *, alignment, block_size +): + """Masks unsafe rows and target positions using device-side invariants.""" + sequence_length = positions.shape[1] + batch_size = positions.shape[0] + validity_mask = jnp.asarray(validity_mask, dtype=jnp.bool_) + completion_mask = jnp.asarray(completion_mask, dtype=jnp.bool_) + corruption_mask = jnp.asarray(corruption_mask, dtype=jnp.bool_) + loss_weights = jnp.asarray(loss_weights) + positions = jnp.asarray(positions, dtype=jnp.int32) + weighted_positions = loss_weights > 0 + completion_is_valid = jnp.all(~completion_mask | validity_mask, axis=1) + corruption_is_completion = jnp.all(~corruption_mask | completion_mask, axis=1) + loss_is_completion = jnp.all(~weighted_positions | completion_mask, axis=1) + corruption_is_weighted = jnp.all(~corruption_mask | weighted_positions, axis=1) + positions_in_range = jnp.all(~validity_mask | ((positions >= 0) & (positions < sequence_length)), axis=1) + valid_completion_mask = completion_mask & validity_mask + valid_corruption_mask = corruption_mask & valid_completion_mask + safe_positions = jnp.clip(positions, 0, sequence_length - 1) + row_indices = jnp.broadcast_to(jnp.arange(batch_size, dtype=jnp.int32)[:, None], positions.shape) + + logical_counts = jnp.zeros_like(positions).at[row_indices, safe_positions].add(validity_mask.astype(jnp.int32)) + valid_count = jnp.sum(validity_mask, axis=1) + expected_counts = jnp.arange(sequence_length, dtype=jnp.int32)[None, :] < valid_count[:, None] + positions_are_canonical = jnp.all(logical_counts == expected_counts.astype(jnp.int32), axis=1) + + logical_completion = jnp.zeros_like(valid_completion_mask).at[row_indices, safe_positions].max(valid_completion_mask) + seen_completion = jax.lax.associative_scan(jnp.logical_or, logical_completion, axis=1) + invalid_future_context = seen_completion & expected_counts & ~logical_completion + has_prompt_origin = expected_counts[:, 0] & ~logical_completion[:, 0] + has_completion = jnp.any(logical_completion, axis=1) + row_is_safe = ( + positions_in_range + & positions_are_canonical + & completion_is_valid + & corruption_is_completion + & loss_is_completion + & corruption_is_weighted + & ~jnp.any(invalid_future_context, axis=1) + & has_prompt_origin + & has_completion + ) + numeric_weights_are_safe = jnp.all(jnp.isfinite(loss_weights) & (loss_weights >= 0), axis=1) + row_is_safe &= numeric_weights_are_safe + + allowed_positions = valid_completion_mask & (positions != 0) + if alignment == "same_position": + allowed_positions &= valid_corruption_mask + else: + shifted_anchors = (positions > 0) & (positions % block_size == 0) + allowed_positions &= valid_corruption_mask | shifted_anchors + row_is_safe &= jnp.all(~weighted_positions | allowed_positions, axis=1) + return jnp.where(row_is_safe[:, None] & allowed_positions, loss_weights, 0.0) + + +def _completion_through_first_stop(generated_tokens, completion_mask, stop_token_ids): + """Keeps completion positions through the first generated stop token, inclusive.""" + if not stop_token_ids: + return completion_mask + stop_positions = completion_mask & jnp.isin(generated_tokens, jnp.asarray(stop_token_ids, generated_tokens.dtype)) + seen_stop = jax.lax.associative_scan(jnp.logical_or, stop_positions, axis=1) + seen_stop_before = jnp.concatenate([jnp.zeros_like(seen_stop[:, :1]), seen_stop[:, :-1]], axis=1) + return completion_mask & ~seen_stop_before + + +def prepare_diffusion_opd_batch( + input_data: distillation_utils.MaxTextTrainingInput, + student_model: nnx.Module, + teacher_model: nnx.Module, + *, + rollout_fn, + teacher_score_fn, + mask_id: int, + stop_token_ids: tuple[int, ...] = (), + shifted_seed: bool = False, + logit_alignment: str = "same_position", + block_size: int = 1, +) -> diffusion_distillation.DiffusionDistillationBatch: + """Creates a fresh rollout and teacher scores for one OPD step.""" + required_fields = { + "targets": input_data.targets, + "positions": input_data.positions, + "completion_mask": input_data.completion_mask, + "corruption_mask": input_data.corruption_mask, + "targets_loss_mask": input_data.targets_loss_mask, + } + missing = sorted(name for name, value in required_fields.items() if value is None) + if missing: + raise ValueError(f"diffusion OPD requires explicit batch fields; missing {missing}") + targets = input_data.targets + positions = input_data.positions + decoder_segment_ids = input_data.decoder_segment_ids + if decoder_segment_ids is None: + decoder_segment_ids = jnp.asarray(input_data.input_mask, dtype=jnp.int32) + validity_mask = decoder_segment_ids != 0 + completion_mask = jnp.asarray(input_data.completion_mask, dtype=jnp.bool_) + corruption_mask = jnp.asarray(input_data.corruption_mask, dtype=jnp.bool_) + loss_weights = jnp.asarray(input_data.targets_loss_mask, dtype=jnp.float32) + _validate_input_masks( + positions, + validity_mask, + completion_mask, + corruption_mask, + loss_weights, + logit_alignment=logit_alignment, + block_size=block_size, + ) + completion_mask &= validity_mask + corruption_mask &= completion_mask + denoise.validate_completion_suffix( + positions, + validity_mask, + completion_mask, + shifted_seed=shifted_seed, + ) + + generated_tokens = rollout_fn( + student_model, + targets, + positions, + decoder_segment_ids, + completion_mask, + ) + completion_mask = _completion_through_first_stop(generated_tokens, completion_mask, stop_token_ids) + validity_mask &= ~jnp.asarray(input_data.completion_mask, dtype=jnp.bool_) | completion_mask + decoder_segment_ids = jnp.where(validity_mask, decoder_segment_ids, 0) + corruption_mask &= completion_mask + loss_weights = jnp.where(completion_mask, loss_weights, 0.0) + student_inputs = jnp.where(corruption_mask, jnp.asarray(mask_id, generated_tokens.dtype), generated_tokens) + loss_weights = _safe_loss_weights( + positions, + validity_mask, + completion_mask, + corruption_mask, + loss_weights, + alignment=logit_alignment, + block_size=block_size, + ) + model_inputs = { + "input_tokens": student_inputs, + "positions": positions, + "decoder_segment_ids": decoder_segment_ids, + "targets": generated_tokens, + "targets_segmentation": decoder_segment_ids, + } + student_batch = diffusion_types.DiffusionTokenBatch.create( + model_inputs=model_inputs, + target_ids=generated_tokens, + loss_weights=loss_weights, + ) + teacher_inputs = dict(model_inputs) + teacher_inputs["input_tokens"] = generated_tokens + teacher_logits = teacher_score_fn(teacher_model, teacher_inputs) + return diffusion_distillation.DiffusionDistillationBatch.create( + student_batch=student_batch, + teacher_logits=teacher_logits, + ) + + +def validate_diffusion_opd_configs(student_config, teacher_config, *, is_offline: bool) -> bool: + """Validates OPD-only constraints and returns whether OPD is enabled.""" + if student_config.distill_data_source != "student_rollout": + return False + if is_offline: + raise ValueError("diffusion OPD requires an online teacher; offline_data_dir is not supported") + if student_config.training_objective != "block_diffusion": + raise ValueError("diffusion OPD requires student training_objective='block_diffusion'") + if not student_config.use_sft or not student_config.sft_train_on_completion_only: + raise ValueError("diffusion OPD requires completion-only SFT data with an explicit prompt/completion boundary") + if teacher_config.training_objective != "causal_lm" or teacher_config.attention_type in { + "block_diffusion", + "full", + "compressed", + }: + raise ValueError("diffusion OPD currently requires a causal teacher to score clean student rollouts") + if getattr(teacher_config, "mtp_num_layers", 0) > 0: + raise ValueError("diffusion OPD does not support an MTP teacher") + if student_config.vocab_size != teacher_config.vocab_size: + raise ValueError("diffusion OPD requires matching student and teacher vocabularies") + student_tokenizer = ( + str(getattr(student_config, "tokenizer_path", "")), + getattr(getattr(student_config, "tokenizer_type", ""), "value", getattr(student_config, "tokenizer_type", "")), + ) + teacher_tokenizer = ( + str(getattr(teacher_config, "tokenizer_path", "")), + getattr(getattr(teacher_config, "tokenizer_type", ""), "value", getattr(teacher_config, "tokenizer_type", "")), + ) + if student_tokenizer != teacher_tokenizer: + raise ValueError( + "diffusion OPD requires identical student and teacher tokenizer_path/tokenizer_type; " + "token ID remapping is not supported" + ) + if student_tokenizer[1] != "huggingface": + raise ValueError("diffusion OPD requires tokenizer_type='huggingface' to match the HF input pipeline") + if student_config.distill_beta != 0.0 or student_config.distill_beta_end is not None: + raise ValueError("diffusion OPD does not support feature distillation") + if student_config.distill_alpha_schedule != "constant" or student_config.distill_temperature_schedule != "constant": + raise ValueError("diffusion OPD currently requires constant alpha and temperature") + if student_config.distill_alpha_end is not None or student_config.distill_temperature_end is not None: + raise ValueError("diffusion OPD does not accept scheduled alpha or temperature end values") + if not math.isfinite(student_config.distill_alpha) or not 0.0 <= student_config.distill_alpha <= 1.0: + raise ValueError("distill_alpha must be finite and in [0, 1] for diffusion OPD") + if not math.isfinite(student_config.distill_temperature) or student_config.distill_temperature <= 0.0: + raise ValueError("distill_temperature must be finite and positive for diffusion OPD") + if not math.isfinite(student_config.distill_rollout_temperature): + raise ValueError("distill_rollout_temperature must be finite for diffusion OPD") + if not math.isfinite(student_config.distill_rollout_confidence_threshold): + raise ValueError("distill_rollout_confidence_threshold must be finite for diffusion OPD") + if student_config.eval_interval > 0: + raise ValueError("diffusion OPD evaluation is not supported yet; set eval_interval=-1") + if student_config.steps <= 0: + raise ValueError("diffusion OPD requires a positive finite number of training steps") + if student_config.gradient_accumulation_steps > 1 and not student_config.use_tunix_gradient_accumulation: + raise ValueError( + "diffusion OPD with gradient accumulation requires use_tunix_gradient_accumulation=True " + "so the HF batch is divided into Tunix microbatches exactly once" + ) + if getattr(student_config, "generate_padding_batch_train", False): + raise ValueError("diffusion OPD does not support generate_padding_batch_train") + if getattr(student_config, "elastic_enabled", False): + raise ValueError("diffusion OPD does not yet support elastic data loading or topology changes") + if getattr(student_config, "expansion_factor_real_data", -1.0) != -1.0: + raise ValueError("diffusion OPD does not support expansion_factor_real_data") + if getattr(student_config, "enable_rampup_batch_size", False): + raise ValueError("diffusion OPD does not support ramp-up batch sizing") + if student_config.learn_to_init_mode: + raise ValueError("diffusion OPD is not compatible with learn_to_init_mode") + if getattr(student_config, "student_params_to_update", None): + raise ValueError("diffusion OPD does not yet support student_params_to_update filtering") + max_steps = student_config.distill_rollout_max_denoise_steps + if max_steps != -1 and max_steps < student_config.block_diffusion_block_size: + raise ValueError("distill_rollout_max_denoise_steps must be -1 or at least block_diffusion_block_size") + stop_token_ids = student_config.distill_rollout_stop_token_ids + if len(set(stop_token_ids)) != len(stop_token_ids) or any( + token_id < 0 or token_id >= student_config.vocab_size for token_id in stop_token_ids + ): + raise ValueError("distill_rollout_stop_token_ids must contain unique IDs in the student vocabulary") + return True + + +def _checkpoint_contract(student_config, teacher_config, stop_token_ids=()): + """Builds the semantic identity that must remain stable across resume.""" + return { + "version": 1, + "student_model": str(getattr(student_config, "model_name", "")), + "student_tokenizer": str(getattr(student_config, "tokenizer_path", "")), + "student_tokenizer_type": str( + getattr(getattr(student_config, "tokenizer_type", ""), "value", getattr(student_config, "tokenizer_type", "")) + ), + "teacher_model": str(getattr(teacher_config, "model_name", "")), + "teacher_checkpoint": str(getattr(teacher_config, "load_parameters_path", "")), + "teacher_tokenizer": str(getattr(teacher_config, "tokenizer_path", "")), + "teacher_tokenizer_type": str( + getattr(getattr(teacher_config, "tokenizer_type", ""), "value", getattr(teacher_config, "tokenizer_type", "")) + ), + "vocab_size": int(student_config.vocab_size), + "block_size": int(student_config.block_diffusion_block_size), + "mask_id": int(student_config.block_diffusion_mask_id), + "stop_token_ids": [int(token_id) for token_id in stop_token_ids], + "min_noise": float(student_config.block_diffusion_min_noise), + "logit_alignment": student_config.block_diffusion_logit_alignment, + "canvas_policy": student_config.block_diffusion_canvas_policy, + "rollout_algorithm": student_config.distill_rollout_algorithm, + "rollout_confidence_threshold": float(student_config.distill_rollout_confidence_threshold), + "rollout_temperature": float(student_config.distill_rollout_temperature), + "rollout_max_denoise_steps": int(student_config.distill_rollout_max_denoise_steps), + "enable_dropout": bool(student_config.enable_dropout), + "dropout_rate": float(getattr(student_config, "dropout_rate", 0.0)), + "steps": int(student_config.steps), + "optimizer_type": str(student_config.opt_type), + "learning_rate": float(student_config.learning_rate), + "warmup_steps_fraction": float(student_config.warmup_steps_fraction), + "learning_rate_final_fraction": float(student_config.learning_rate_final_fraction), + "gradient_clipping_threshold": float(student_config.gradient_clipping_threshold), + "adam_b1": float(getattr(student_config, "adam_b1", 0.0)), + "adam_b2": float(getattr(student_config, "adam_b2", 0.0)), + "adam_eps": float(getattr(student_config, "adam_eps", 0.0)), + "adam_eps_root": float(getattr(student_config, "adam_eps_root", 0.0)), + "adam_weight_decay": float(getattr(student_config, "adam_weight_decay", 0.0)), + "adamw_mask": str(getattr(student_config, "adamw_mask", "")), + "mu_dtype": str(getattr(student_config, "mu_dtype", "")), + "skip_step_on_spikes": bool(getattr(student_config, "skip_step_on_spikes", False)), + "skip_step_interval": int(getattr(student_config, "skip_step_interval", 0)), + "skip_step_scaling_factor": float(getattr(student_config, "skip_step_scaling_factor", 0.0)), + "trainable_parameters_mask": str(getattr(student_config, "trainable_parameters_mask", "")), + "distill_temperature": float(student_config.distill_temperature), + "soft_loss_weight": float(student_config.distill_alpha), + "hard_loss_weight": float(1.0 - student_config.distill_alpha), + "dataset_type": str(student_config.dataset_type), + "hf_path": str(getattr(student_config, "hf_path", "")), + "hf_name": str(getattr(student_config, "hf_name", "")), + "hf_data_dir": str(getattr(student_config, "hf_data_dir", "")), + "hf_train_files": str(getattr(student_config, "hf_train_files", "")), + "train_split": str(getattr(student_config, "train_split", "")), + "train_data_columns": str(getattr(student_config, "train_data_columns", "")), + "tokenize_train_data": bool(getattr(student_config, "tokenize_train_data", True)), + "enable_data_shuffling": bool(getattr(student_config, "enable_data_shuffling", False)), + "data_shuffle_seed": int(getattr(student_config, "data_shuffle_seed", 0)), + "num_epoch": int(getattr(student_config, "num_epoch", 1)), + "max_target_length": int(student_config.max_target_length), + "add_bos": bool(getattr(student_config, "add_bos", True)), + "add_eos": bool(getattr(student_config, "add_eos", True)), + "packing": bool(getattr(student_config, "packing", False)), + "generate_padding_batch_train": bool(getattr(student_config, "generate_padding_batch_train", False)), + "chat_template_path": str(getattr(student_config, "chat_template_path", "")), + "chat_template": str(getattr(student_config, "chat_template", "")), + "formatting_func_path": str(getattr(student_config, "formatting_func_path", "")), + "formatting_func_kwargs": str(getattr(student_config, "formatting_func_kwargs", "")), + "gradient_accumulation_steps": int(student_config.gradient_accumulation_steps), + "use_tunix_gradient_accumulation": bool(student_config.use_tunix_gradient_accumulation), + "global_batch_size_to_load": int(student_config.global_batch_size_to_load), + "global_batch_size_to_train_on": int(getattr(student_config, "global_batch_size_to_train_on", 0)), + "data_sharding": str(getattr(student_config, "data_sharding", "")), + "expansion_factor_real_data": float(getattr(student_config, "expansion_factor_real_data", -1.0)), + "process_count": int(jax.process_count()), + "local_device_count": int(jax.local_device_count()), + } + + +def replay_and_bound_iterator( + iterator, + *, + iter_steps: int, + train_steps: int, + max_steps: int, + accumulation_steps: int, + replay_iterator=None, +): + """Replays a deterministic HF stream and bounds it to remaining updates. + + `replay_iterator` may be the host-local source underlying `iterator`. Advancing + it avoids materializing and device-putting global arrays for skipped batches. + """ + if iter_steps < 0 or train_steps < 0 or accumulation_steps < 1: + raise ValueError("restored iterator and training step counts must be nonnegative") + if train_steps > max_steps: + raise ValueError(f"restored train_steps ({train_steps}) exceeds configured max_steps ({max_steps})") + if iter_steps != train_steps * accumulation_steps: + raise ValueError("restored iter_steps must equal train_steps * accumulation_steps") + remaining_updates = max_steps - train_steps + if remaining_updates == 0: + return iter(()) + replay_source = iterator if replay_iterator is None else replay_iterator + for replayed_steps in range(iter_steps): + try: + next(replay_source) + except StopIteration as exc: + raise ValueError( + f"diffusion OPD HF replay exhausted after {replayed_steps} of {iter_steps} required microbatches" + ) from exc + remaining_microbatches = remaining_updates * accumulation_steps + + def take_remaining_exactly(): + for consumed_steps in range(remaining_microbatches): + try: + yield next(iterator) + except StopIteration as exc: + raise ValueError( + "diffusion OPD HF stream exhausted after " + f"{consumed_steps} of {remaining_microbatches} required remaining microbatches" + ) from exc + + return take_remaining_exactly() + + +class MaxTextDiffusionOPDTrainer(peft_trainer.PeftTrainer): + """Tunix trainer that prepares fresh MaxText rollouts before every JIT step.""" + + def __init__( + self, + *, + student_model, + teacher_model, + optimizer, + training_config, + student_config, + teacher_config, + eos_id, + ): + super().__init__(model=student_model, optimizer=optimizer, training_config=training_config) + if teacher_model is None: + raise ValueError("diffusion OPD requires a teacher model") + configured_stop_ids = tuple(student_config.distill_rollout_stop_token_ids) + if not configured_stop_ids and (eos_id is None or eos_id < 0 or eos_id >= student_config.vocab_size): + raise ValueError("diffusion OPD requires a valid tokenizer EOS token ID") + self.teacher_model = teacher_model + self.student_config = student_config + self.stop_token_ids = configured_stop_ids or (int(eos_id),) + self.checkpoint_contract = _checkpoint_contract(student_config, teacher_config, self.stop_token_ids) + self.rollout_fn = create_rollout_fn(student_config) + self.teacher_score_fn = create_teacher_score_fn(teacher_config) + student_logits_fn = create_target_aligned_logits_fn(student_config, enable_dropout=student_config.enable_dropout) + tunix_diffusion_opd.configure_prepared_diffusion_opd( + self, + lambda batch: batch, + student_logits_fn, + temperature=student_config.distill_temperature, + soft_loss_weight=student_config.distill_alpha, + hard_loss_weight=1.0 - student_config.distill_alpha, + ) + + def _prepare_inputs(self, input_data): + return prepare_diffusion_opd_batch( + input_data, + self.model, + self.teacher_model, + rollout_fn=self.rollout_fn, + teacher_score_fn=self.teacher_score_fn, + mask_id=self.student_config.block_diffusion_mask_id, + stop_token_ids=self.stop_token_ids, + shifted_seed=self.student_config.block_diffusion_logit_alignment == "shifted", + logit_alignment=self.student_config.block_diffusion_logit_alignment, + block_size=self.student_config.block_diffusion_block_size, + ) + + def setup_checkpoint_manager_and_restore(self, raw_train_iter, config): + """Restores model state; HF stream position is replayed deterministically.""" + if self.checkpoint_manager is not None: + self.checkpoint_manager.close() + self.checkpoint_manager = distillation_utils.MaxTextCheckpointManager( + raw_iterator=None, + root_directory=config.checkpoint_dir, + student_config=config, + options=self.config.checkpointing_options, + ) + self._train_steps, self._restored_custom_metadata = self.checkpoint_manager.maybe_restore( + self.model, + self.optimizer, + restore_only_lora_params=getattr(self, "_lora_enabled", False), + ) + if self._train_steps > 0: + restored_contract = self._restored_custom_metadata.get("diffusion_opd_contract") + if restored_contract is None: + raise ValueError("diffusion OPD checkpoint is missing its semantic training contract") + if restored_contract != self.checkpoint_contract: + raise ValueError( + "diffusion OPD checkpoint contract does not match the current student, teacher, tokenizer, or objective" + ) + self._iter_steps = self._train_steps * self.config.get_with_default("gradient_accumulation_steps", 1) + if self._iter_steps: + max_logging.log(f"Diffusion OPD will replay {self._iter_steps} deterministic HF microbatches after restore.") + return raw_train_iter + + def custom_checkpoint_metadata(self): + """Returns the immutable semantic contract persisted with each checkpoint.""" + return {"diffusion_opd_contract": self.checkpoint_contract} diff --git a/src/maxtext/trainers/post_train/distillation/distillation_utils.py b/src/maxtext/trainers/post_train/distillation/distillation_utils.py index cfa7fdfeb2..e8e9bba025 100644 --- a/src/maxtext/trainers/post_train/distillation/distillation_utils.py +++ b/src/maxtext/trainers/post_train/distillation/distillation_utils.py @@ -69,6 +69,12 @@ class MaxTextTrainingInput(peft_trainer.TrainingInput): targets_position: jax.Array | None = None #: Segment IDs for packed target tokens. targets_segmentation: jax.Array | None = None + #: Role-derived positions eligible for completion generation. + completion_mask: jax.Array | None = None + #: Positions replaced by the diffusion mask token for the student forward pass. + corruption_mask: jax.Array | None = None + #: Per-target weights for target-aligned diffusion objectives. + targets_loss_mask: jax.Array | None = None #: Top-K logits from the teacher model. top_k_logits: jax.Array | None = None top_k_indices: jax.Array | None = None @@ -132,6 +138,9 @@ def __next__(self) -> MaxTextTrainingInput: targets=batch["targets"], targets_position=targets_position, targets_segmentation=targets_segmentation, + completion_mask=batch.get("completion_mask"), + corruption_mask=batch.get("corruption_mask"), + targets_loss_mask=batch.get("targets_loss_mask"), top_k_logits=batch.get("top_k_logits"), top_k_indices=batch.get("top_k_indices"), ) diff --git a/src/maxtext/trainers/post_train/distillation/train_distill.py b/src/maxtext/trainers/post_train/distillation/train_distill.py index f31b43b921..84ae819a3c 100644 --- a/src/maxtext/trainers/post_train/distillation/train_distill.py +++ b/src/maxtext/trainers/post_train/distillation/train_distill.py @@ -422,6 +422,9 @@ def _prepare_inputs( targets=input_data.targets, targets_position=input_data.targets_position, targets_segmentation=input_data.targets_segmentation, + completion_mask=input_data.completion_mask, + corruption_mask=input_data.corruption_mask, + targets_loss_mask=input_data.targets_loss_mask, top_k_logits=input_data.top_k_logits, top_k_indices=input_data.top_k_indices, ) @@ -532,6 +535,11 @@ def setup_checkpoint_manager_and_restore(self, raw_train_iter, config): self.optimizer, restore_only_lora_params=getattr(self, "_lora_enabled", False), ) + if "diffusion_opd_contract" in (self._restored_custom_metadata or {}): + raise ValueError( + "standard distillation cannot resume a diffusion OPD checkpoint; " + "keep distill_data_source='student_rollout' or start from a compatible checkpoint" + ) grad_accum_steps = self.config.get_with_default("gradient_accumulation_steps", 1) self._iter_steps = self._train_steps * grad_accum_steps @@ -574,14 +582,14 @@ def build_training_components( is_offline: bool = False, offline_data_dir: str | None = None, ): - """Builds and returns the strategy, optimizer, and training config objects. + """Builds the strategy, optimizer, training config, and tokenizer EOS ID. Args: student_config: Configuration object for the Student model. teacher_config: Configuration object for the Teacher model. Returns: - A tuple of (DistillationStrategy, Optimizer, TrainingConfig). + A tuple of (DistillationStrategy, Optimizer, TrainingConfig, eos_id). """ # 2. Load Tokenizer Info tok = tokenizer.build_tokenizer( @@ -642,6 +650,9 @@ def build_training_components( log_dir=student_config.tensorboard_dir, flush_every_n_steps=student_config.log_period ) + opd_training_options = ( + {"max_inflight_computations": 1} if student_config.distill_data_source == "student_rollout" else {} + ) train_config = peft_trainer.TrainingConfig( max_steps=student_config.steps, eval_every_n_steps=student_config.eval_interval, @@ -651,9 +662,10 @@ def build_training_components( checkpointing_options=checkpointing_options, gradient_accumulation_steps=student_config.gradient_accumulation_steps, data_sharding_axis=tuple(student_config.data_sharding), + **opd_training_options, ) - return strategy, optimizer, train_config + return strategy, optimizer, train_config, tok.eos_id def train_distill( @@ -677,9 +689,19 @@ def train_distill( f"Vocab size mismatch! Student: {student_config.vocab_size}, Teacher: {teacher_config.vocab_size}. " "Distillation requires matching vocabularies." ) + diffusion_opd_enabled = student_config.distill_data_source == "student_rollout" + diffusion_opd_lib = None + if diffusion_opd_enabled: + # Keep the standard distillation import path compatible with the currently + # pinned Tunix release. New Tunix diffusion APIs are required only when OPD + # is explicitly selected. + # pylint: disable-next=import-outside-toplevel + from maxtext.trainers.post_train.distillation import diffusion_opd as diffusion_opd_lib + + diffusion_opd_lib.validate_diffusion_opd_configs(student_config, teacher_config, is_offline=is_offline) # Build Training Components (No hardware context required) - strategy, optimizer, train_config = build_training_components( + strategy, optimizer, train_config, eos_id = build_training_components( student_config, teacher_config, is_offline, offline_data_dir ) @@ -724,21 +746,33 @@ def student_freeze_param_fn(path) -> bool: ) student_model.train() - model_bundle = ModelBundle(teacher_model, student_model) - - # 3. Initialize Trainer - trainer = MaxTextDistillationTrainer( - model=model_bundle, - strategy=strategy, - optimizer=optimizer, - training_config=train_config, - student_config=student_config, - teacher_config=teacher_config, - is_offline=is_offline, - student_freeze_param_filter=student_freeze_param_fn if student_params_to_update else None, - ) + model_bundle = None + if diffusion_opd_enabled: + assert diffusion_opd_lib is not None + trainer = diffusion_opd_lib.MaxTextDiffusionOPDTrainer( + student_model=student_model, + teacher_model=teacher_model, + optimizer=optimizer, + training_config=train_config, + student_config=student_config, + teacher_config=teacher_config, + eos_id=eos_id, + ) + else: + model_bundle = ModelBundle(teacher_model, student_model) + trainer = MaxTextDistillationTrainer( + model=model_bundle, + strategy=strategy, + optimizer=optimizer, + training_config=train_config, + student_config=student_config, + teacher_config=teacher_config, + is_offline=is_offline, + student_freeze_param_filter=student_freeze_param_fn if student_params_to_update else None, + ) trainer.is_managed_externally = True - trainer._has_aux = True # pylint: disable=protected-access + if not diffusion_opd_enabled: + trainer._has_aux = True # pylint: disable=protected-access if is_offline: max_logging.log("Initializing Data Iterators via MaxText pipeline...") @@ -759,7 +793,10 @@ def student_freeze_param_fn(path) -> bool: # Sync the ModelBundle step counter with the restored training step so that # loss weight schedules resume from the correct position after checkpoint restore. - model_bundle.training_step.set_value(jnp.array(trainer._train_steps, dtype=jnp.int32)) # pylint: disable=protected-access + if model_bundle is not None: + model_bundle.training_step.set_value( # pylint: disable=protected-access + jnp.array(trainer._train_steps, dtype=jnp.int32) # pylint: disable=protected-access + ) # 6. Configure Input Mapping def custom_gen_model_input_fn(batch): @@ -782,10 +819,21 @@ def custom_gen_model_input_fn(batch): ) return inputs_dict - trainer = trainer.with_gen_model_input_fn(custom_gen_model_input_fn) + if not diffusion_opd_enabled: + trainer = trainer.with_gen_model_input_fn(custom_gen_model_input_fn) # 7. Create Iterator Wrappers (Use Utils) train_iter = distillation_utils.MaxTextToTunixIterator(raw_train_iter) + if diffusion_opd_enabled: + assert diffusion_opd_lib is not None + train_iter = diffusion_opd_lib.replay_and_bound_iterator( + train_iter, + iter_steps=trainer.iter_steps, + train_steps=trainer.train_steps, + max_steps=student_config.steps, + accumulation_steps=student_config.gradient_accumulation_steps, + replay_iterator=getattr(raw_train_iter, "local_iterator", None), + ) eval_iter = None if raw_eval_iter is not None: @@ -821,6 +869,7 @@ def custom_gen_model_input_fn(batch): optimizer=trainer.optimizer, save_only_lora_params=getattr(trainer, "_lora_enabled", False), force=True, + custom_metadata=trainer.custom_checkpoint_metadata() if diffusion_opd_enabled else {}, ) if saved: # Ensure underlying orbax manager finishes writing @@ -911,6 +960,12 @@ def main(argv: Sequence[str]) -> None: teacher_argv = [argv[0], argv[1]] teacher_config = pyconfig.initialize(teacher_argv, **teacher_overrides) + if student_config.distill_data_source == "student_rollout": + # pylint: disable-next=import-outside-toplevel + from maxtext.trainers.post_train.distillation import diffusion_opd as diffusion_opd_lib + + diffusion_opd_lib.validate_diffusion_opd_configs(student_config, teacher_config, is_offline=is_offline) + # Batch shape (per_device_batch_size / max_target_length / gradient_accumulation_steps) # must be set at the YAML top level — not inside *_overrides — since student and # teacher share the input pipeline. diff --git a/src/maxtext/trainers/post_train/sft/hooks.py b/src/maxtext/trainers/post_train/sft/hooks.py index 9f30a3d6d8..03edfbcdb9 100644 --- a/src/maxtext/trainers/post_train/sft/hooks.py +++ b/src/maxtext/trainers/post_train/sft/hooks.py @@ -29,7 +29,7 @@ class SFTTrainingHooks(BaseTrainingHooks): @override def get_total_weights(self, batch) -> jax.Array: """Calculate the number of non-padded tokens in the batch.""" - return jnp.sum(batch["targets_segmentation"] != 0) + return jnp.sum(batch.get("targets_loss_mask", batch["targets_segmentation"]) != 0) class SFTDataHooks(BaseDataHooks): diff --git a/src/maxtext/trainers/post_train/sft/train_sft.py b/src/maxtext/trainers/post_train/sft/train_sft.py index 9002caa16c..715585ec90 100644 --- a/src/maxtext/trainers/post_train/sft/train_sft.py +++ b/src/maxtext/trainers/post_train/sft/train_sft.py @@ -49,10 +49,12 @@ from orbax import checkpoint as ocp +from tunix.sft import diffusion as tunix_diffusion_sft from tunix.sft import metrics_logger, peft_trainer, profiler from maxtext.optimizers import optimizers from maxtext.configs import pyconfig +from maxtext.integration.tunix import diffusion_sft as maxtext_diffusion_sft from maxtext.trainers.pre_train.train import loss_fn from maxtext.common.goodput import ( GoodputEvent, @@ -210,6 +212,9 @@ def loss_func( targets, targets_position, targets_segmentation, + completion_mask=None, + corruption_mask=None, + targets_loss_mask=None, ): data = { "inputs": inputs, @@ -219,12 +224,36 @@ def loss_func( "targets_position": targets_position, "targets_segmentation": targets_segmentation, } + if completion_mask is not None: + data["completion_mask"] = completion_mask + if corruption_mask is not None: + data["corruption_mask"] = corruption_mask + if targets_loss_mask is not None: + data["targets_loss_mask"] = targets_loss_mask return loss_fn(model, mt_config, data, dropout_rng=None, params=None, is_train=True) trainer = trainer.with_loss_fn(loss_func, has_aux=True) return trainer +def configure_training_objective(trainer, mt_config): + """Configures AR or target-aligned diffusion loss without changing AR defaults.""" + if getattr(mt_config, "training_objective", "causal_lm") != "block_diffusion": + return use_maxtext_loss_function(trainer, mt_config) + max_logging.log("Configuring Tunix target-aligned block-diffusion SFT adapter.") + return tunix_diffusion_sft.configure_diffusion_sft( + trainer, + maxtext_diffusion_sft.create_batch_adapter(mt_config), + maxtext_diffusion_sft.create_target_aligned_logits_fn(mt_config), + ) + + +def _create_trainer(model, optimizer, tunix_config, mt_config): + if getattr(mt_config, "training_objective", "causal_lm") == "block_diffusion": + return peft_trainer.PeftTrainer(model, optimizer, tunix_config) + return MaxTextPeftTrainer(model, optimizer, tunix_config) + + def validate_config(config): """Validates the configuration parameters for SFT training.""" if config.optimizer_memory_host_offload: @@ -237,6 +266,7 @@ def validate_config(config): def setup_trainer_state(mt_config, goodput_recorder=None): """Set up prerequisites for training loop.""" + validate_config(mt_config) tunix_config = get_tunix_config(mt_config) with maybe_record_goodput(goodput_recorder, GoodputEvent.TPU_INIT): @@ -261,12 +291,12 @@ def setup_trainer_state(mt_config, goodput_recorder=None): # Provide rules context so 'norm' is translated to mesh axes during maybe_restore with nn_partitioning.axis_rules(mt_config.logical_axis_rules): - trainer = MaxTextPeftTrainer(model, optimizer, tunix_config) + trainer = _create_trainer(model, optimizer, tunix_config, mt_config) if mt_config.lora.lora_restore_path and trainer.train_steps == 0: lora_utils.restore_lora_from_path(trainer.model, mt_config) trainer.with_training_hooks(training_hooks) trainer.with_data_hooks(data_hooks) - trainer = use_maxtext_loss_function(trainer, mt_config) + trainer = configure_training_objective(trainer, mt_config) return trainer, mesh diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index ee92b93435..f27addd189 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -41,6 +41,7 @@ from flax.nnx import variablelib from maxtext.configs import pyconfig +from maxtext.diffusion import scoring as diffusion_scoring from maxtext.utils.globals import EPS from maxtext.utils import elastic_utils # Placeholder: internal @@ -103,6 +104,22 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr loss: average loss aux: a dictionary including intermediate_outputs, xent_sum, and total_weights """ + is_block_diffusion = getattr(config, "training_objective", "causal_lm") == "block_diffusion" + if getattr(config, "attention_type", "global") == "block_diffusion" and not is_block_diffusion: + raise ValueError( + "Block-diffusion attention requires target-aligned block-diffusion losses; " + "causal next-token labels would leak within a bidirectional block." + ) + if is_block_diffusion: + required_masks = {"completion_mask", "corruption_mask", "targets_loss_mask"} + missing_masks = required_masks - data.keys() + if missing_masks: + raise ValueError(f"Block-diffusion loss requires explicit batch masks; missing {sorted(missing_masks)}") + target_shape = data["targets"].shape + for mask_name in required_masks: + if data[mask_name].shape != target_shape: + raise ValueError(f"{mask_name} must match targets shape; got {data[mask_name].shape} and {target_shape}") + # decimate proportion of data when per_device_batch_size<1 if is_train: for k, v in data.items(): @@ -110,6 +127,14 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr else: for k, v in data.items(): data[k] = v[: config.micro_batch_size_to_eval_on, :] + loss_mask_values = data["targets_loss_mask"] if is_block_diffusion else data["targets_segmentation"] + targets_loss_mask = (loss_mask_values != 0) & (data["targets_segmentation"] != 0) + target_positions = data.get("targets_position", data["inputs_position"]) + if is_block_diffusion and config.block_diffusion_canvas_policy == "seed_and_mask": + if config.block_diffusion_logit_alignment == "shifted": + targets_loss_mask &= target_positions != 0 + else: + targets_loss_mask &= target_positions % config.block_diffusion_block_size != 0 mutable_collections = ["intermediates"] if config.mtp_num_layers > 0 and is_train: # The single model.apply call now triggers the entire chain if MTP is enabled: @@ -161,6 +186,13 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0] xent_sum, total_z_loss = vocab_tiling_linen_loss(hidden_states, data, config, model, params, is_train) else: + if is_block_diffusion: + logits = diffusion_scoring.align_logits_to_targets( + logits, + config.block_diffusion_logit_alignment, + target_positions, + data["targets_segmentation"] != 0, + ) one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size) xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier) @@ -179,9 +211,8 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr debug_sharding=config.debug_sharding, ) - # Mask out paddings at the end of each example. - xent = xent * (data["targets_segmentation"] != 0) - z_loss = z_loss * (data["targets_segmentation"] != 0) + xent = xent * targets_loss_mask + z_loss = z_loss * targets_loss_mask xent_sum = jnp.sum(xent) total_z_loss = jnp.sum(z_loss) @@ -224,6 +255,13 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0] xent_sum, total_z_loss = vocab_tiling_nnx_loss(model, hidden_states, data, config, is_train) else: + if is_block_diffusion: + logits = diffusion_scoring.align_logits_to_targets( + logits, + config.block_diffusion_logit_alignment, + target_positions, + data["targets_segmentation"] != 0, + ) one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size) xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier) @@ -242,14 +280,13 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr debug_sharding=config.debug_sharding, ) - # Mask out paddings at the end of each example. - xent = xent * (data["targets_segmentation"] != 0) - z_loss = z_loss * (data["targets_segmentation"] != 0) + xent = xent * targets_loss_mask + z_loss = z_loss * targets_loss_mask xent_sum = jnp.sum(xent) total_z_loss = jnp.sum(z_loss) - total_weights = jnp.sum(data["targets_segmentation"] != 0) + total_weights = jnp.sum(targets_loss_mask) # If gradient accumulation is enabled, we don't need to divide xent_sum # by total_weights and then multiply the computed gradient by total_weights, # since it's equivalent to computing the gradient from xent_sum. diff --git a/src/maxtext/utils/gradient_accumulation.py b/src/maxtext/utils/gradient_accumulation.py index 35fdf65503..eb0786c812 100644 --- a/src/maxtext/utils/gradient_accumulation.py +++ b/src/maxtext/utils/gradient_accumulation.py @@ -162,12 +162,15 @@ def reshape_to_microbatch_accumulations(batch_arr): grad_and_loss, aux = jax.lax.scan( accumulate_gradient, init_grad_and_loss, data, length=config.gradient_accumulation_steps ) + has_weights = grad_and_loss["total_weights"] > 0 + safe_total_weights = jnp.maximum(grad_and_loss["total_weights"], 1) loss = ( - grad_and_loss["loss"] / grad_and_loss["total_weights"] + grad_and_loss["loss"] / safe_total_weights + grad_and_loss["moe_lb_loss"] / config.gradient_accumulation_steps + grad_and_loss["indexer_loss"] / config.gradient_accumulation_steps + grad_and_loss["mtp_loss"] / config.gradient_accumulation_steps ) + loss = jnp.where(has_weights, loss, 0.0) raw_grads = grad_and_loss["grad"] if data_parallel_active: # Mark the gradients unreduced over the "data" axis now that we're outside the @@ -176,7 +179,10 @@ def reshape_to_microbatch_accumulations(batch_arr): unreduced_shardings = jax.tree.map(update_sharding_for_unreduced, params_shardings) raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, unreduced_shardings) raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, params_shardings) - raw_grads = jax.tree_util.tree_map(lambda arr: arr / grad_and_loss["total_weights"], raw_grads) + raw_grads = jax.tree_util.tree_map( + lambda arr: jnp.where(has_weights, arr / safe_total_weights, jnp.zeros_like(arr)), + raw_grads, + ) aux = jax.tree.map(lambda x: jnp.sum(x, axis=0), aux) # pytype: disable=module-attr if is_nnx: diff --git a/src/maxtext/utils/max_utils.py b/src/maxtext/utils/max_utils.py index e6aebce059..943b1c63c2 100644 --- a/src/maxtext/utils/max_utils.py +++ b/src/maxtext/utils/max_utils.py @@ -957,6 +957,9 @@ def reorder_causal_load_balanced(batch, cp_size, reorder_strategy, hardware="tpu "targets_position", "inputs_segmentation", "targets_segmentation", + "completion_mask", + "corruption_mask", + "targets_loss_mask", } if hardware in ("gpu", "gpu_multiprocess"): diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 6c71e0487e..7f5845a72d 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -161,6 +161,10 @@ def get_shaped_batch(config, batch_sharding=None): shaped_batch["targets"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) shaped_batch["targets_position"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) shaped_batch["targets_segmentation"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) + if getattr(config, "training_objective", "causal_lm") == "block_diffusion": + shaped_batch["completion_mask"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) + shaped_batch["corruption_mask"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) + shaped_batch["targets_loss_mask"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) if config.use_multimodal: image_shape = mm_processor.get_dummy_image_shape_for_init( config.model_name, batch_size=config.micro_batch_size_to_train_on diff --git a/tests/post_training/unit/diffusion_rl_test.py b/tests/post_training/unit/diffusion_rl_test.py new file mode 100644 index 0000000000..62ce3b7abe --- /dev/null +++ b/tests/post_training/unit/diffusion_rl_test.py @@ -0,0 +1,470 @@ +# Copyright 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 +# +# http://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. + +"""Tests exact MaxText denoising-trace replay for diffusion RL.""" + +from types import SimpleNamespace + +from absl.testing import absltest +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.diffusion import denoise +from maxtext.diffusion import scoring +from maxtext.integration.tunix import diffusion_rl +from tunix.rl import diffusion as tunix_diffusion + + +def _config(**overrides): + """Builds the minimal diffusion configuration used by these tests.""" + values = { + "block_diffusion_block_size": 4, + "block_diffusion_mask_id": 7, + "block_diffusion_logit_alignment": "same_position", + "block_diffusion_canvas_policy": "all_masked", + "vocab_size": 8, + "data_shuffle_seed": 3, + "rl": SimpleNamespace( + diffusion_stop_token_ids=[], + diffusion_max_denoise_steps=-1, + diffusion_confidence_threshold=0.99, + ), + } + values.update(overrides) + return SimpleNamespace(**values) + + +class _CanvasDependentModel(nnx.Module): + + def __init__(self): + self.scale = nnx.Param(jnp.asarray(2.0, dtype=jnp.float32)) + + def __call__(self, decoder_input_tokens, **kwargs): + del kwargs + committed = jnp.sum(decoder_input_tokens != 7, axis=1, keepdims=True) + target_ids = (committed + jnp.arange(decoder_input_tokens.shape[1])[None, :]) % 7 + return jax.nn.one_hot(target_ids, 8, dtype=jnp.float32) * self.scale[...] + + +class _NonFiniteModel(nnx.Module): + + def __init__(self): + self.scale = nnx.Param(jnp.asarray(1.0, dtype=jnp.float32)) + + def __call__(self, decoder_input_tokens, **kwargs): + del kwargs + return jnp.full((*decoder_input_tokens.shape, 8), jnp.nan) * self.scale[...] + + +class _Tokenizer: + """Minimal tokenizer contract used by the in-process rollout tests.""" + + def encode(self, text): + del text + return [4] + + def decode(self, token_ids): + return " ".join(str(token_id) for token_id in token_ids) + + def pad_id(self): + return 0 + + def eos_id(self): + return 6 + + +class _PadValuedPromptTokenizer(_Tokenizer): + + def encode(self, text): + del text + return [0, 4] + + +class DiffusionRLTest(absltest.TestCase): + + def test_trace_replay_matches_action_time_logps_and_has_live_gradient(self): + model = _CanvasDependentModel() + initial = jnp.asarray([[4, 1, 1, 1]], dtype=jnp.int32) + positions = jnp.arange(4, dtype=jnp.int32)[None, :] + completion_mask = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_) + temperature = 0.7 + + trace = denoise.low_confidence_rollout( + lambda canvas: model(decoder_input_tokens=canvas), + initial, + positions, + jnp.ones_like(completion_mask), + completion_mask, + block_size=4, + mask_id=7, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + temperature=temperature, + rng=jax.random.PRNGKey(7), + ) + batch = diffusion_rl.prepare_diffusion_policy_batch( + prompt_tokens=initial[:, :1], + prompt_mask=jnp.ones((1, 1), dtype=jnp.bool_), + completion_tokens=trace.tokens[:, 1:], + completion_mask=trace.action_steps[:, 1:] >= 0, + action_steps=trace.action_steps[:, 1:], + ) + logits_fn = diffusion_rl.make_diffusion_trace_logits_fn(_config()) + replayed = tunix_diffusion.compute_diffusion_per_token_logps( + model, + batch, + logits_fn, + temperature=temperature, + stop_gradient=False, + ) + + np.testing.assert_allclose(replayed, trace.action_logps[:, 1:], rtol=1e-6, atol=1e-6) + gradients = nnx.grad( + lambda scored_model: jnp.sum( + tunix_diffusion.compute_diffusion_per_token_logps( + scored_model, + batch, + logits_fn, + temperature=temperature, + stop_gradient=False, + ) + ) + )(model) + self.assertGreater(float(jnp.abs(gradients.scale[...])), 0.0) + + def test_shifted_seed_trace_replay_matches_action_time_logps(self): + model = _CanvasDependentModel() + initial = jnp.asarray([[2, 3, 4, 5, 7, 7, 7, 7]], dtype=jnp.int32) + positions = jnp.arange(8, dtype=jnp.int32)[None, :] + completion_mask = jnp.asarray([[0, 0, 0, 0, 1, 1, 1, 1]], dtype=jnp.bool_) + temperature = 0.7 + + trace = denoise.low_confidence_rollout( + lambda canvas: scoring.align_logits_to_targets( + model(decoder_input_tokens=canvas), + "shifted", + positions, + jnp.ones_like(completion_mask), + ), + initial, + positions, + jnp.ones_like(completion_mask), + completion_mask, + block_size=4, + mask_id=7, + logit_alignment="shifted", + canvas_policy="seed_and_mask", + confidence_threshold=0.99, + temperature=temperature, + rng=jax.random.PRNGKey(9), + ) + batch = diffusion_rl.prepare_diffusion_policy_batch( + prompt_tokens=initial[:, :4], + prompt_mask=jnp.ones((1, 4), dtype=jnp.bool_), + completion_tokens=trace.tokens[:, 4:], + completion_mask=jnp.ones((1, 4), dtype=jnp.bool_), + action_steps=trace.action_steps[:, 4:], + ) + config = _config( + block_diffusion_logit_alignment="shifted", + block_diffusion_canvas_policy="seed_and_mask", + ) + + replayed = tunix_diffusion.compute_diffusion_per_token_logps( + model, + batch, + diffusion_rl.make_diffusion_trace_logits_fn(config), + temperature=temperature, + stop_gradient=True, + ) + + np.testing.assert_allclose(replayed, trace.action_logps[:, 4:], rtol=1e-6, atol=1e-6) + + def test_prepared_batch_rejects_action_mask_mismatch(self): + with self.assertRaisesRegex(ValueError, "exactly one nonnegative action step"): + diffusion_rl.prepare_diffusion_policy_batch( + prompt_tokens=jnp.asarray([[4]]), + prompt_mask=jnp.asarray([[1]], dtype=jnp.bool_), + completion_tokens=jnp.asarray([[5, 6]]), + completion_mask=jnp.asarray([[1, 1]], dtype=jnp.bool_), + action_steps=jnp.asarray([[0, -1]]), + ) + + def test_stop_tokens_mask_actions_after_first_match(self): + keep = diffusion_rl.policy_mask_at_stop( + jnp.asarray([[4, 5, 6, 9]]), + jnp.asarray([[0, 1, 2, 3]]), + (5, 8), + ) + + np.testing.assert_array_equal(keep, [[1, 1, 0, 0]]) + + def test_stop_mask_preserves_full_trace_for_exact_replay(self): + model = _CanvasDependentModel() + initial = jnp.asarray([[4, 7, 7, 7]], dtype=jnp.int32) + positions = jnp.arange(4, dtype=jnp.int32)[None, :] + completion_mask = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_) + temperature = 0.7 + trace = denoise.low_confidence_rollout( + lambda canvas: model(decoder_input_tokens=canvas), + initial, + positions, + jnp.ones_like(completion_mask), + completion_mask, + block_size=4, + mask_id=7, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + temperature=temperature, + rng=jax.random.PRNGKey(7), + ) + completion_tokens = trace.tokens[:, 1:] + completion_steps = trace.action_steps[:, 1:] + stop_id = int(completion_tokens[0, 0]) + policy_mask = diffusion_rl.policy_mask_at_stop(completion_tokens, completion_steps, (stop_id,)) + batch = diffusion_rl.prepare_diffusion_policy_batch( + prompt_tokens=initial[:, :1], + prompt_mask=jnp.ones((1, 1), dtype=jnp.bool_), + completion_tokens=completion_tokens, + completion_mask=jnp.ones_like(completion_tokens, dtype=jnp.bool_), + action_steps=completion_steps, + loss_mask=policy_mask, + inactive_target_id=0, + ) + + replayed = tunix_diffusion.compute_diffusion_per_token_logps( + model, + batch, + diffusion_rl.make_diffusion_trace_logits_fn(_config()), + temperature=temperature, + stop_gradient=True, + ) + + np.testing.assert_array_equal(batch.model_inputs["action_steps"], completion_steps) + np.testing.assert_array_equal(batch.target_ids, jnp.where(policy_mask, completion_tokens, 0)) + np.testing.assert_allclose( + replayed, + jnp.where(policy_mask, trace.action_logps[:, 1:], 0.0), + rtol=1e-6, + atol=1e-6, + ) + + def test_explicit_stop_ids_do_not_require_tokenizer_eos(self): + config = _config( + rl=SimpleNamespace( + diffusion_stop_token_ids=[5, 6], + diffusion_max_denoise_steps=-1, + diffusion_confidence_threshold=0.99, + ) + ) + + self.assertEqual(diffusion_rl.resolve_stop_token_ids(config, None), (5, 6)) + + def test_prompt_mask_comes_from_padding_length_not_token_value(self): + rollout_config = SimpleNamespace(temperature=0.7) + rollout = diffusion_rl.MaxTextDiffusionRollout( + rollout_actor=_CanvasDependentModel(), + tokenizer=_PadValuedPromptTokenizer(), + mesh=None, + rollout_config=rollout_config, + maxtext_config=_config(), + ) + + tokens, mask = rollout._encode_left_padded(["prompt"], 3) # pylint: disable=protected-access + + np.testing.assert_array_equal(tokens, [[0, 0, 4]]) + np.testing.assert_array_equal(mask, [[0, 1, 1]]) + + def test_rollout_exports_padded_trace_and_matching_sampler_logps(self): + config = _config() + model = _CanvasDependentModel() + rollout_config = SimpleNamespace( + max_prompt_length=2, + max_tokens_to_generate=3, + temperature=0.7, + return_logprobs=True, + top_k=-1, + top_p=1.0, + ) + rollout = diffusion_rl.MaxTextDiffusionRollout( + rollout_actor=model, + tokenizer=_Tokenizer(), + mesh=None, + rollout_config=rollout_config, + maxtext_config=config, + ) + + output = rollout.generate(["prompt", "prompt"], rollout_config) + batch = output.diffusion_batch + self.assertIsNotNone(batch) + self.assertIsInstance(batch.target_ids, np.ndarray) + self.assertEqual(batch.target_ids.shape, (2, 3)) + np.testing.assert_array_equal(batch.loss_weights > 0, batch.model_inputs["action_steps"] >= 0) + np.testing.assert_array_equal(batch.target_ids, batch.model_inputs["completion_tokens"]) + replayed = tunix_diffusion.compute_diffusion_per_token_logps( + model, + batch, + diffusion_rl.make_diffusion_trace_logits_fn(config), + temperature=rollout_config.temperature, + stop_gradient=True, + ) + np.testing.assert_allclose(replayed, np.stack(output.logprobs), rtol=1e-6, atol=1e-6) + + def test_rollout_scores_full_trace_while_reward_text_stops_at_eos(self): + config = _config( + rl=SimpleNamespace( + diffusion_stop_token_ids=list(range(7)), + diffusion_max_denoise_steps=-1, + diffusion_confidence_threshold=0.99, + ) + ) + training_config = SimpleNamespace( + max_prompt_length=2, + max_tokens_to_generate=3, + temperature=0.7, + return_logprobs=True, + top_k=-1, + top_p=1.0, + ) + rollout = diffusion_rl.MaxTextDiffusionRollout( + rollout_actor=_CanvasDependentModel(), + tokenizer=_Tokenizer(), + mesh=None, + rollout_config=training_config, + maxtext_config=config, + ) + + output = rollout.generate(["prompt"], training_config) + + self.assertLen(output.tokens[0], 1) + self.assertLen(output.logprobs[0], 3) + np.testing.assert_array_equal(output.diffusion_batch.loss_weights, np.ones((1, 3))) + + def test_greedy_eval_uses_its_own_sampling_contract(self): + training_config = SimpleNamespace(temperature=0.7) + rollout = diffusion_rl.MaxTextDiffusionRollout( + rollout_actor=_CanvasDependentModel(), + tokenizer=_Tokenizer(), + mesh=None, + rollout_config=training_config, + maxtext_config=_config(), + ) + eval_config = SimpleNamespace( + max_prompt_length=2, + max_tokens_to_generate=3, + temperature=0.01, + return_logprobs=False, + top_k=1, + top_p=1.0, + ) + + first = rollout.generate(["prompt"], eval_config) + second = rollout.generate(["prompt"], eval_config) + + np.testing.assert_array_equal(first.diffusion_batch.target_ids, second.diffusion_batch.target_ids) + + def test_rollout_fails_loudly_on_non_finite_policy_scores(self): + rollout_config = SimpleNamespace( + max_prompt_length=2, + max_tokens_to_generate=3, + temperature=0.7, + return_logprobs=True, + top_k=-1, + top_p=1.0, + ) + rollout = diffusion_rl.MaxTextDiffusionRollout( + rollout_actor=_NonFiniteModel(), + tokenizer=_Tokenizer(), + mesh=None, + rollout_config=rollout_config, + maxtext_config=_config(), + ) + + with self.assertRaisesRegex(RuntimeError, "unresolved mask|non-finite"): + rollout.generate(["prompt"], rollout_config) + + def test_eval_rng_does_not_advance_training_stream(self): + training_config = SimpleNamespace( + max_prompt_length=2, + max_tokens_to_generate=3, + temperature=0.7, + return_logprobs=True, + top_k=-1, + top_p=1.0, + seed=13, + ) + eval_config = SimpleNamespace( + max_prompt_length=2, + max_tokens_to_generate=3, + temperature=0.01, + return_logprobs=False, + top_k=1, + top_p=1.0, + ) + + def make_rollout(): + return diffusion_rl.MaxTextDiffusionRollout( + rollout_actor=_CanvasDependentModel(), + tokenizer=_Tokenizer(), + mesh=None, + rollout_config=training_config, + maxtext_config=_config(), + ) + + with_eval = make_rollout() + baseline = make_rollout() + with_eval.generate(["prompt"], eval_config) + + actual = with_eval.generate(["prompt"], training_config) + expected = baseline.generate(["prompt"], training_config) + + np.testing.assert_array_equal(actual.diffusion_batch.target_ids, expected.diffusion_batch.target_ids) + + def test_generation_context_restores_training_rng_by_global_step(self): + training_config = SimpleNamespace( + max_prompt_length=2, + max_tokens_to_generate=3, + temperature=0.7, + return_logprobs=True, + top_k=-1, + top_p=1.0, + seed=13, + ) + + def make_rollout(): + return diffusion_rl.MaxTextDiffusionRollout( + rollout_actor=_CanvasDependentModel(), + tokenizer=_Tokenizer(), + mesh=None, + rollout_config=training_config, + maxtext_config=_config(), + ) + + uninterrupted = make_rollout() + resumed = make_rollout() + uninterrupted.set_generation_context(global_step=7, mode="train") + resumed.set_generation_context(global_step=7, mode="train") + + expected = uninterrupted.generate(["prompt"], training_config) + actual = resumed.generate(["prompt"], training_config) + + np.testing.assert_array_equal(actual.diffusion_batch.target_ids, expected.diffusion_batch.target_ids) + + +if __name__ == "__main__": + absltest.main() diff --git a/tests/post_training/unit/diffusion_sft_test.py b/tests/post_training/unit/diffusion_sft_test.py new file mode 100644 index 0000000000..475b0ecdd7 --- /dev/null +++ b/tests/post_training/unit/diffusion_sft_test.py @@ -0,0 +1,95 @@ +# Copyright 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 +# +# http://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. + +"""Tests for the MaxText-to-Tunix diffusion SFT adapter.""" + +from types import SimpleNamespace + +import jax.numpy as jnp +import numpy as np +import pytest + +from maxtext.diffusion import scoring +from maxtext.integration.tunix import diffusion_sft + + +pytestmark = [pytest.mark.post_training, pytest.mark.cpu_only] + + +def _config(alignment="shifted"): + return SimpleNamespace( + block_diffusion_block_size=4, + block_diffusion_logit_alignment=alignment, + enable_dropout=False, + ) + + +def _raw_batch(): + positions = jnp.arange(8, dtype=jnp.int32)[None, :] + segmentation = jnp.ones((1, 8), dtype=jnp.int32) + return { + "inputs": jnp.asarray([[10, 11, 99, 99, 14, 99, 99, 99]], dtype=jnp.int32), + "inputs_position": positions, + "inputs_segmentation": segmentation, + "targets": jnp.asarray([[10, 11, 12, 13, 14, 15, 16, 17]], dtype=jnp.int32), + "targets_position": positions, + "targets_segmentation": segmentation, + "completion_mask": jnp.asarray([[0, 0, 1, 1, 1, 1, 1, 1]], dtype=jnp.int32), + "corruption_mask": jnp.asarray([[0, 0, 1, 1, 0, 1, 1, 1]], dtype=jnp.int32), + "targets_loss_mask": jnp.asarray([[0, 0, 1, 1, 1, 1, 1, 1]], dtype=jnp.int32), + } + + +def test_batch_adapter_preserves_shifted_anchor_weight(): + batch = diffusion_sft.create_batch_adapter(_config())(_raw_batch()) + + np.testing.assert_array_equal(batch.target_ids, _raw_batch()["targets"]) + np.testing.assert_array_equal(batch.loss_weights, [[0, 0, 1, 1, 1, 1, 1, 1]]) + assert batch.loss_weights.dtype == jnp.float32 + + +def test_batch_adapter_rejects_unowned_clean_target(): + raw = _raw_batch() + raw["corruption_mask"] = raw["corruption_mask"].at[0, 3].set(0) + + with pytest.raises(ValueError, match="corrupted targets or shifted block anchors"): + diffusion_sft.create_batch_adapter(_config())(raw) + + +def test_logits_adapter_uses_target_alignment(): + raw_logits = jnp.arange(1 * 8 * 3, dtype=jnp.float32).reshape(1, 8, 3) + calls = [] + + class Model: + + def __call__(self, **kwargs): + calls.append(kwargs) + return raw_logits + + batch = diffusion_sft.create_batch_adapter(_config())(_raw_batch()) + actual = diffusion_sft.create_target_aligned_logits_fn(_config())( + Model(), batch.model_inputs + ) + expected = scoring.align_logits_to_targets( + raw_logits, + "shifted", + batch.model_inputs["target_positions"], + batch.model_inputs["target_segmentation"] != 0, + ) + + np.testing.assert_array_equal(actual, expected) + assert calls[0]["enable_dropout"] is False + np.testing.assert_array_equal( + calls[0]["decoder_input_tokens"], _raw_batch()["inputs"] + ) diff --git a/tests/post_training/unit/maxtext_diffusion_opd_test.py b/tests/post_training/unit/maxtext_diffusion_opd_test.py new file mode 100644 index 0000000000..96ae731155 --- /dev/null +++ b/tests/post_training/unit/maxtext_diffusion_opd_test.py @@ -0,0 +1,539 @@ +# Copyright 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 +# +# http://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. + +"""Tests MaxText preparation of fresh diffusion OPD batches.""" + +from types import SimpleNamespace +from unittest import mock + +from absl.testing import absltest +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.trainers.post_train.distillation import diffusion_opd +from maxtext.trainers.post_train.distillation import distillation_utils + + +def _student_config(**overrides): + """Builds the student configuration subset used by OPD validation.""" + values = { + "distill_data_source": "student_rollout", + "training_objective": "block_diffusion", + "attention_type": "block_diffusion", + "use_sft": True, + "sft_train_on_completion_only": True, + "vocab_size": 16, + "tokenizer_path": "test-tokenizer", + "tokenizer_type": "huggingface", + "distill_beta": 0.0, + "distill_beta_end": None, + "distill_alpha": 0.75, + "distill_alpha_end": None, + "distill_alpha_schedule": "constant", + "distill_temperature": 2.0, + "distill_temperature_end": None, + "distill_temperature_schedule": "constant", + "eval_interval": -1, + "steps": 10, + "gradient_accumulation_steps": 1, + "use_tunix_gradient_accumulation": False, + "max_target_length": 16, + "global_batch_size_to_load": 8, + "global_batch_size_to_train_on": 8, + "dataset_type": "hf", + "generate_padding_batch_train": False, + "elastic_enabled": False, + "expansion_factor_real_data": -1.0, + "enable_rampup_batch_size": False, + "learn_to_init_mode": False, + "student_params_to_update": None, + "distill_rollout_max_denoise_steps": -1, + "block_diffusion_block_size": 4, + "block_diffusion_mask_id": 15, + "block_diffusion_min_noise": 0.001, + "block_diffusion_logit_alignment": "same_position", + "block_diffusion_canvas_policy": "all_masked", + "distill_rollout_confidence_threshold": 0.9, + "distill_rollout_temperature": 1.0, + "distill_rollout_algorithm": "low_confidence", + "distill_rollout_stop_token_ids": [], + "enable_dropout": True, + "dropout_rate": 0.0, + "opt_type": "adamw", + "learning_rate": 1e-4, + "warmup_steps_fraction": 0.1, + "learning_rate_final_fraction": 0.1, + "gradient_clipping_threshold": 1.0, + "hardware": "cpu", + "shard_mode": "auto", + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _teacher_config(**overrides): + """Builds a causal teacher configuration subset.""" + values = { + "training_objective": "causal_lm", + "attention_type": "global", + "vocab_size": 16, + "tokenizer_path": "test-tokenizer", + "tokenizer_type": "huggingface", + "hardware": "cpu", + "shard_mode": "auto", + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _training_input(**overrides): + """Builds one explicit-mask target-aligned distillation batch.""" + values = { + "input_tokens": jnp.asarray([[3, 4, 31, 31]], dtype=jnp.int32), + "input_mask": jnp.ones((1, 4), dtype=jnp.bool_), + "positions": jnp.arange(4, dtype=jnp.int32)[None, :], + "decoder_segment_ids": jnp.ones((1, 4), dtype=jnp.int32), + "targets": jnp.asarray([[3, 4, 5, 6]], dtype=jnp.int32), + "targets_position": jnp.arange(4, dtype=jnp.int32)[None, :], + "targets_segmentation": jnp.ones((1, 4), dtype=jnp.int32), + "completion_mask": jnp.asarray([[0, 0, 1, 1]], dtype=jnp.int32), + "corruption_mask": jnp.asarray([[0, 0, 1, 1]], dtype=jnp.int32), + "targets_loss_mask": jnp.asarray([[0, 0, 1, 1]], dtype=jnp.int32), + } + values.update(overrides) + return distillation_utils.MaxTextTrainingInput(**values) + + +class MaxTextDiffusionOPDTest(absltest.TestCase): + + def test_config_validation_is_default_off(self): + student = _student_config(distill_data_source="dataset") + + self.assertFalse(diffusion_opd.validate_diffusion_opd_configs(student, _teacher_config(), is_offline=False)) + + def test_config_validation_accepts_supported_contract(self): + self.assertTrue(diffusion_opd.validate_diffusion_opd_configs(_student_config(), _teacher_config(), is_offline=False)) + + def test_config_validation_rejects_unsafe_modes(self): + cases = [ + (_student_config(), _teacher_config(), True, "online teacher"), + (_student_config(training_objective="causal_lm"), _teacher_config(), False, "student training_objective"), + (_student_config(use_sft=False), _teacher_config(), False, "completion-only SFT"), + (_student_config(), _teacher_config(training_objective="block_diffusion"), False, "causal teacher"), + (_student_config(), _teacher_config(attention_type="full"), False, "causal teacher"), + (_student_config(), _teacher_config(mtp_num_layers=1), False, "MTP teacher"), + (_student_config(), _teacher_config(tokenizer_path="other-tokenizer"), False, "identical.*tokenizer"), + ( + _student_config(tokenizer_type="sentencepiece"), + _teacher_config(tokenizer_type="sentencepiece"), + False, + "tokenizer_type='huggingface'", + ), + (_student_config(distill_beta=0.1), _teacher_config(), False, "feature distillation"), + (_student_config(distill_alpha=float("nan")), _teacher_config(), False, "finite"), + (_student_config(distill_rollout_temperature=float("inf")), _teacher_config(), False, "finite"), + (_student_config(eval_interval=10), _teacher_config(), False, "evaluation"), + ( + _student_config(generate_padding_batch_train=True), + _teacher_config(), + False, + "generate_padding_batch_train", + ), + (_student_config(elastic_enabled=True), _teacher_config(), False, "elastic data loading"), + (_student_config(expansion_factor_real_data=2.0), _teacher_config(), False, "expansion_factor_real_data"), + (_student_config(enable_rampup_batch_size=True), _teacher_config(), False, "ramp-up batch sizing"), + ( + _student_config(gradient_accumulation_steps=2), + _teacher_config(), + False, + "use_tunix_gradient_accumulation=True", + ), + (_student_config(distill_rollout_max_denoise_steps=3), _teacher_config(), False, "at least"), + ( + _student_config(distill_rollout_stop_token_ids=[7, 7]), + _teacher_config(), + False, + "unique IDs", + ), + ] + for student, teacher, is_offline, message in cases: + with self.subTest(message=message), self.assertRaisesRegex(ValueError, message): + diffusion_opd.validate_diffusion_opd_configs(student, teacher, is_offline=is_offline) + + def test_tpu_auto_sharding_is_preserved_for_models_without_explicit_support(self): + self.assertTrue( + diffusion_opd.validate_diffusion_opd_configs( + _student_config(hardware="tpu", shard_mode="auto"), + _teacher_config(hardware="tpu", shard_mode="auto"), + is_offline=False, + ) + ) + + def test_causal_teacher_logits_are_shifted_to_physical_targets(self): + raw_logits = jnp.arange(8, dtype=jnp.float32).reshape(1, 4, 2) + model = mock.Mock(return_value=raw_logits) + logits_fn = diffusion_opd.create_target_aligned_logits_fn(_teacher_config(), enable_dropout=False) + model_inputs = { + "input_tokens": jnp.ones((1, 4), dtype=jnp.int32), + "positions": jnp.arange(4, dtype=jnp.int32)[None, :], + "decoder_segment_ids": jnp.ones((1, 4), dtype=jnp.int32), + "targets": jnp.ones((1, 4), dtype=jnp.int32), + "targets_segmentation": jnp.ones((1, 4), dtype=jnp.int32), + } + + aligned = logits_fn(model, model_inputs) + + np.testing.assert_array_equal(aligned, raw_logits[:, [0, 0, 1, 2], :]) + self.assertFalse(model.call_args.kwargs["enable_dropout"]) + + def test_compiled_rollout_functionalizes_mutating_nnx_model(self): + class MutatingLogitModel(nnx.Module): + + def __init__(self): + self.counter = nnx.Variable(jnp.asarray(0, dtype=jnp.int32)) + + def __call__(self, decoder_input_tokens, **kwargs): + del kwargs + self.counter[...] += 1 + targets = jnp.full_like(decoder_input_tokens, 7) + return jax.nn.one_hot(targets, 16, dtype=jnp.float32) * 12.0 + + model = MutatingLogitModel() + rollout_fn = diffusion_opd.create_rollout_fn(_student_config()) + + generated = rollout_fn( + model, + jnp.asarray([[3, 4, 1, 1]], dtype=jnp.int32), + jnp.arange(4, dtype=jnp.int32)[None, :], + jnp.ones((1, 4), dtype=jnp.int32), + jnp.asarray([[0, 0, 1, 1]], dtype=jnp.bool_), + ) + + np.testing.assert_array_equal(generated, [[3, 4, 7, 7]]) + self.assertEqual(int(model.counter[...]), 0) + + def test_preparation_scores_clean_rollout_and_corrupts_only_student(self): + generated = jnp.asarray([[3, 4, 7, 8]], dtype=jnp.int32) + rollout_fn = mock.Mock(return_value=generated) + teacher_model = mock.sentinel.teacher_model + + def teacher_score_fn(model, model_inputs): + self.assertIs(model, teacher_model) + np.testing.assert_array_equal(model_inputs["input_tokens"], generated) + return jax.nn.one_hot(generated, 16, dtype=jnp.float32) + + batch = diffusion_opd.prepare_diffusion_opd_batch( + _training_input(), + mock.sentinel.student_model, + teacher_model, + rollout_fn=rollout_fn, + teacher_score_fn=teacher_score_fn, + mask_id=15, + ) + + np.testing.assert_array_equal(batch.student_batch.target_ids, generated) + np.testing.assert_array_equal(batch.student_batch.model_inputs["input_tokens"], [[3, 4, 15, 15]]) + np.testing.assert_array_equal(batch.student_batch.loss_weights, [[0.0, 0.0, 1.0, 1.0]]) + self.assertEqual(batch.teacher_logits.shape, (1, 4, 16)) + rollout_fn.assert_called_once() + + def test_preparation_excludes_tokens_after_generated_eos(self): + generated = jnp.asarray([[3, 4, 7, 8]], dtype=jnp.int32) + teacher_score_fn = mock.Mock(return_value=jax.nn.one_hot(generated, 16)) + + batch = diffusion_opd.prepare_diffusion_opd_batch( + _training_input(), + mock.sentinel.student_model, + mock.sentinel.teacher_model, + rollout_fn=mock.Mock(return_value=generated), + teacher_score_fn=teacher_score_fn, + mask_id=15, + stop_token_ids=(7,), + ) + + np.testing.assert_array_equal(batch.student_batch.loss_weights, [[0.0, 0.0, 1.0, 0.0]]) + np.testing.assert_array_equal(batch.student_batch.model_inputs["decoder_segment_ids"], [[1, 1, 1, 0]]) + np.testing.assert_array_equal(teacher_score_fn.call_args.args[1]["targets_segmentation"], [[1, 1, 1, 0]]) + + def test_shifted_block_anchor_can_be_weighted_without_corruption(self): + generated = jnp.asarray([[3, 7, 8, 9]], dtype=jnp.int32) + batch = diffusion_opd.prepare_diffusion_opd_batch( + _training_input( + completion_mask=jnp.asarray([[0, 1, 1, 1]], dtype=jnp.int32), + corruption_mask=jnp.asarray([[0, 1, 0, 1]], dtype=jnp.int32), + targets_loss_mask=jnp.asarray([[0, 1, 1, 1]], dtype=jnp.int32), + ), + mock.sentinel.student, + mock.sentinel.teacher, + rollout_fn=mock.Mock(return_value=generated), + teacher_score_fn=mock.Mock(return_value=jax.nn.one_hot(generated, 16)), + mask_id=15, + logit_alignment="shifted", + block_size=2, + ) + + np.testing.assert_array_equal(batch.student_batch.model_inputs["input_tokens"], [[3, 15, 8, 15]]) + np.testing.assert_array_equal(batch.student_batch.loss_weights, [[0.0, 1.0, 1.0, 1.0]]) + + def test_prepared_batch_runs_weighted_tunix_loss_and_student_gradient(self): + class TrainableLogitModel(nnx.Module): + + def __init__(self): + self.bias = nnx.Param(jnp.zeros((16,), dtype=jnp.float32)) + + def __call__(self, decoder_input_tokens, **kwargs): + del kwargs + return jnp.broadcast_to(self.bias[...], (*decoder_input_tokens.shape, 16)) + + generated = jnp.asarray([[3, 4, 7, 8]], dtype=jnp.int32) + batch = diffusion_opd.prepare_diffusion_opd_batch( + _training_input(), + mock.sentinel.student, + mock.sentinel.teacher, + rollout_fn=mock.Mock(return_value=generated), + teacher_score_fn=mock.Mock(return_value=jax.nn.one_hot(generated, 16) * 4.0), + mask_id=15, + ) + model = TrainableLogitModel() + logits_fn = diffusion_opd.create_target_aligned_logits_fn(_student_config(), enable_dropout=False) + + def loss_fn(trainable_model): + output = diffusion_opd.tunix_diffusion_opd.diffusion_opd_loss_fn( + trainable_model, + batch, + logits_fn, + temperature=2.0, + soft_loss_weight=0.75, + hard_loss_weight=0.25, + ) + return output.primary_loss.unreduced_sum, output + + (loss_sum, output), gradients = nnx.value_and_grad(loss_fn, has_aux=True)(model) + + self.assertTrue(bool(jnp.isfinite(loss_sum))) + self.assertEqual(float(output.primary_loss.denominator), 2.0) + self.assertTrue(any(bool(jnp.any(leaf[...] != 0)) for leaf in jax.tree.leaves(gradients))) + + def test_same_position_rejects_weighted_clean_targets(self): + with self.assertRaisesRegex(ValueError, "same-position targets must be corrupted"): + diffusion_opd.prepare_diffusion_opd_batch( + _training_input( + corruption_mask=jnp.asarray([[0, 0, 1, 0]], dtype=jnp.int32), + targets_loss_mask=jnp.asarray([[0, 0, 1, 1]], dtype=jnp.int32), + ), + mock.sentinel.student, + mock.sentinel.teacher, + rollout_fn=mock.Mock(), + teacher_score_fn=mock.Mock(), + mask_id=15, + logit_alignment="same_position", + block_size=4, + ) + + def test_device_safe_weights_zero_invalid_suffix_rows(self): + weights = diffusion_opd._safe_loss_weights( # pylint: disable=protected-access + jnp.arange(4, dtype=jnp.int32)[None, :], + jnp.ones((1, 4), dtype=jnp.bool_), + jnp.asarray([[0, 1, 0, 1]], dtype=jnp.bool_), + jnp.asarray([[0, 1, 0, 1]], dtype=jnp.bool_), + jnp.asarray([[0.0, 1.0, 0.0, 1.0]], dtype=jnp.float32), + alignment="same_position", + block_size=4, + ) + + np.testing.assert_array_equal(weights, jnp.zeros((1, 4), dtype=jnp.float32)) + + def test_device_safe_weights_zero_out_of_range_positions(self): + for positions in (jnp.asarray([[-1, 1]]), jnp.asarray([[0, 2]])): + with self.subTest(positions=positions): + weights = diffusion_opd._safe_loss_weights( # pylint: disable=protected-access + positions, + jnp.ones((1, 2), dtype=jnp.bool_), + jnp.asarray([[0, 1]], dtype=jnp.bool_), + jnp.asarray([[0, 1]], dtype=jnp.bool_), + jnp.asarray([[0.0, 1.0]], dtype=jnp.float32), + alignment="same_position", + block_size=2, + ) + + np.testing.assert_array_equal(weights, jnp.zeros((1, 2), dtype=jnp.float32)) + + def test_device_safe_weights_zero_mask_subset_violations(self): + weights = diffusion_opd._safe_loss_weights( # pylint: disable=protected-access + jnp.asarray([[0, 1, 2]]), + jnp.asarray([[1, 1, 0]], dtype=jnp.bool_), + jnp.asarray([[0, 1, 1]], dtype=jnp.bool_), + jnp.asarray([[0, 1, 1]], dtype=jnp.bool_), + jnp.asarray([[0.0, 1.0, 1.0]], dtype=jnp.float32), + alignment="same_position", + block_size=2, + ) + + np.testing.assert_array_equal(weights, jnp.zeros((1, 3), dtype=jnp.float32)) + + def test_preparation_rejects_missing_or_cross_role_masks(self): + rollout_fn = mock.Mock() + teacher_score_fn = mock.Mock() + with self.assertRaisesRegex(ValueError, "corruption_mask"): + diffusion_opd.prepare_diffusion_opd_batch( + _training_input(corruption_mask=None), + mock.sentinel.student, + mock.sentinel.teacher, + rollout_fn=rollout_fn, + teacher_score_fn=teacher_score_fn, + mask_id=15, + ) + with self.assertRaisesRegex(ValueError, "subset of completion_mask"): + diffusion_opd.prepare_diffusion_opd_batch( + _training_input(corruption_mask=jnp.asarray([[0, 1, 1, 0]], dtype=jnp.int32)), + mock.sentinel.student, + mock.sentinel.teacher, + rollout_fn=rollout_fn, + teacher_score_fn=teacher_score_fn, + mask_id=15, + ) + rollout_fn.assert_not_called() + teacher_score_fn.assert_not_called() + + def test_hf_resume_replays_and_bounds_deterministic_stream(self): + bounded = diffusion_opd.replay_and_bound_iterator( + iter(range(20)), + iter_steps=4, + train_steps=2, + max_steps=5, + accumulation_steps=2, + ) + + self.assertEqual(list(bounded), [4, 5, 6, 7, 8, 9]) + + def test_hf_resume_can_replay_before_global_materialization(self): + local_iterator = iter(range(20)) + + class GlobalIterator: + + def __next__(self): + return next(local_iterator) * 10 + + def __iter__(self): + return self + + bounded = diffusion_opd.replay_and_bound_iterator( + GlobalIterator(), + iter_steps=4, + train_steps=2, + max_steps=5, + accumulation_steps=2, + replay_iterator=local_iterator, + ) + + self.assertEqual(list(bounded), [40, 50, 60, 70, 80, 90]) + + def test_hf_resume_rejects_early_exhaustion_during_replay(self): + with self.assertRaisesRegex(ValueError, "replay exhausted after 2 of 4"): + diffusion_opd.replay_and_bound_iterator( + iter(range(2)), + iter_steps=4, + train_steps=2, + max_steps=5, + accumulation_steps=2, + ) + + def test_hf_resume_rejects_early_exhaustion_during_remaining_training(self): + bounded = diffusion_opd.replay_and_bound_iterator( + iter(range(4)), + iter_steps=4, + train_steps=2, + max_steps=5, + accumulation_steps=2, + replay_iterator=iter(range(4)), + ) + + with self.assertRaisesRegex(ValueError, "4 of 6 required remaining microbatches"): + list(bounded) + + def test_hf_resume_rejects_checkpoint_beyond_configured_steps(self): + with self.assertRaisesRegex(ValueError, "exceeds configured max_steps"): + diffusion_opd.replay_and_bound_iterator( + iter(()), + iter_steps=12, + train_steps=6, + max_steps=5, + accumulation_steps=2, + ) + + def test_hf_resume_at_final_step_does_not_replay(self): + bounded = diffusion_opd.replay_and_bound_iterator( + iter(()), + iter_steps=10, + train_steps=5, + max_steps=5, + accumulation_steps=2, + ) + + self.assertEqual(list(bounded), []) + + def test_checkpoint_contract_captures_batch_and_preprocessing_identity(self): + base = diffusion_opd._checkpoint_contract(_student_config(), _teacher_config()) # pylint: disable=protected-access + changed = diffusion_opd._checkpoint_contract( # pylint: disable=protected-access + _student_config(gradient_accumulation_steps=2, use_tunix_gradient_accumulation=True), _teacher_config() + ) + + self.assertNotEqual(base, changed) + self.assertEqual(base["max_target_length"], 16) + + def test_checkpoint_restore_preserves_contract_without_claiming_hf_cursor_state(self): + trainer = diffusion_opd.MaxTextDiffusionOPDTrainer.__new__(diffusion_opd.MaxTextDiffusionOPDTrainer) + trainer.checkpoint_manager = mock.Mock() + trainer.model = mock.sentinel.model + trainer.optimizer = mock.sentinel.optimizer + trainer.config = mock.Mock(checkpointing_options=mock.sentinel.options) + trainer.config.get_with_default.return_value = 1 + trainer._lora_enabled = False # pylint: disable=protected-access + trainer.checkpoint_contract = {"version": 1, "teacher_model": "teacher"} + manager = mock.Mock() + manager.maybe_restore.return_value = (4, {"diffusion_opd_contract": trainer.checkpoint_contract}) + raw_iterator = SimpleNamespace(local_iterator=mock.sentinel.local_iterator) + + with mock.patch.object(diffusion_opd.distillation_utils, "MaxTextCheckpointManager", return_value=manager) as factory: + result = trainer.setup_checkpoint_manager_and_restore( + raw_iterator, SimpleNamespace(dataset_type="hf", checkpoint_dir="/tmp") + ) + + self.assertIs(result, raw_iterator) + self.assertIsNone(factory.call_args.kwargs["raw_iterator"]) + self.assertEqual(trainer._train_steps, 4) # pylint: disable=protected-access + self.assertEqual(trainer.custom_checkpoint_metadata()["diffusion_opd_contract"], trainer.checkpoint_contract) + + def test_checkpoint_restore_rejects_semantic_contract_mismatch(self): + trainer = diffusion_opd.MaxTextDiffusionOPDTrainer.__new__(diffusion_opd.MaxTextDiffusionOPDTrainer) + trainer.checkpoint_manager = mock.Mock() + trainer.model = mock.sentinel.model + trainer.optimizer = mock.sentinel.optimizer + trainer.config = mock.Mock(checkpointing_options=mock.sentinel.options) + trainer.config.get_with_default.return_value = 1 + trainer._lora_enabled = False # pylint: disable=protected-access + trainer.checkpoint_contract = {"version": 1, "teacher_model": "expected"} + manager = mock.Mock() + manager.maybe_restore.return_value = (4, {"diffusion_opd_contract": {"version": 1, "teacher_model": "other"}}) + + with ( + mock.patch.object(diffusion_opd.distillation_utils, "MaxTextCheckpointManager", return_value=manager), + self.assertRaisesRegex(ValueError, "checkpoint contract"), + ): + trainer.setup_checkpoint_manager_and_restore( + SimpleNamespace(local_iterator=mock.sentinel.local_iterator), + SimpleNamespace(dataset_type="hf", checkpoint_dir="/tmp"), + ) diff --git a/tests/post_training/unit/sft_hooks_test.py b/tests/post_training/unit/sft_hooks_test.py index 9fde5add7e..8b97aceaa5 100644 --- a/tests/post_training/unit/sft_hooks_test.py +++ b/tests/post_training/unit/sft_hooks_test.py @@ -74,6 +74,18 @@ def test_sft_training_hooks_get_total_weights(self): total_weights = training_hooks.get_total_weights(batch) self.assertEqual(total_weights, 3) + def test_sft_training_hooks_prefers_explicit_loss_mask(self): + learning_rate_schedule = maxtext_utils.create_learning_rate_schedule(self.config) + training_hooks = sft_hooks.SFTTrainingHooks(self.config, self.mesh, learning_rate_schedule, goodput_recorder=None) + batch = { + "targets_segmentation": np.ones((2, 4), dtype=np.int32), + "targets_loss_mask": np.array([[1, 0, 0, 0], [0, 1, 1, 0]], dtype=np.int32), + } + + total_weights = training_hooks.get_total_weights(batch) + + self.assertEqual(total_weights, 3) + if __name__ == "__main__": unittest.main() diff --git a/tests/post_training/unit/train_distill_test.py b/tests/post_training/unit/train_distill_test.py index 24cd1d0bc8..6a13c80ef0 100644 --- a/tests/post_training/unit/train_distill_test.py +++ b/tests/post_training/unit/train_distill_test.py @@ -104,6 +104,9 @@ def test_maxtext_to_tunix_iterator_sft(self): "targets": np.array([[11, 12]]), "targets_position": np.array([[100, 101]]), # Custom position "targets_segmentation": np.array([[0, 1]]), # Custom segmentation (mask) + "completion_mask": np.array([[0, 1]]), + "corruption_mask": np.array([[0, 1]]), + "targets_loss_mask": np.array([[0, 1]]), } dummy_iter_sft = iter([dummy_batch_sft]) @@ -115,6 +118,9 @@ def test_maxtext_to_tunix_iterator_sft(self): self.assertIsInstance(tunix_input_sft, distillation_utils.MaxTextTrainingInput) np.testing.assert_array_equal(tunix_input_sft.targets_position, dummy_batch_sft["targets_position"]) np.testing.assert_array_equal(tunix_input_sft.targets_segmentation, dummy_batch_sft["targets_segmentation"]) + np.testing.assert_array_equal(tunix_input_sft.completion_mask, dummy_batch_sft["completion_mask"]) + np.testing.assert_array_equal(tunix_input_sft.corruption_mask, dummy_batch_sft["corruption_mask"]) + np.testing.assert_array_equal(tunix_input_sft.targets_loss_mask, dummy_batch_sft["targets_loss_mask"]) def test_maxtext_to_tunix_iterator_packed_fallback(self): """Verifies fallback behavior when segmentation is missing.""" @@ -151,13 +157,19 @@ def test_prepare_inputs_logic(self): input_mask=jnp.array([[True]]), positions=jnp.array([[0]]), targets=jnp.array([[1]]), + completion_mask=jnp.array([[1]]), + corruption_mask=jnp.array([[1]]), + targets_loss_mask=jnp.array([[1]]), ) # 4. Run - _ = trainer._prepare_inputs(input_data) + prepared = trainer._prepare_inputs(input_data) # 5. Verify trainer.strategy.get_teacher_outputs.assert_not_called() + np.testing.assert_array_equal(prepared.completion_mask, input_data.completion_mask) + np.testing.assert_array_equal(prepared.corruption_mask, input_data.corruption_mask) + np.testing.assert_array_equal(prepared.targets_loss_mask, input_data.targets_loss_mask) @mock.patch("maxtext.trainers.post_train.distillation.train_distill.optax.global_norm") @mock.patch("maxtext.trainers.post_train.distillation.train_distill.jax.tree.map") @@ -646,6 +658,32 @@ def test_setup_pipeline_restored(self): # Verify it returned the restored iterator, NOT the raw one self.assertEqual(result, restored_iter) + def test_standard_distillation_rejects_diffusion_opd_checkpoint(self): + mock_trainer = mock.Mock() + mock_trainer.checkpoint_manager = mock.Mock() + mock_trainer.model = mock.Mock() + mock_trainer.optimizer = mock.Mock() + mock_trainer._lora_enabled = False + mock_trainer.config = mock.Mock( + checkpointing_options=ocp.CheckpointManagerOptions(max_to_keep=1, create=True) + ) + mock_manager = mock.Mock() + mock_manager.maybe_restore.return_value = (10, {"diffusion_opd_contract": {"version": 1}}) + config = mock.Mock(dataset_type="synthetic", checkpoint_dir=self.test_dir) + + with ( + mock.patch( + "maxtext.trainers.post_train.distillation.distillation_utils.MaxTextCheckpointManager", + return_value=mock_manager, + ), + self.assertRaisesRegex(ValueError, "cannot resume a diffusion OPD checkpoint"), + ): + train_distill.MaxTextDistillationTrainer.setup_checkpoint_manager_and_restore( + mock_trainer, + object(), + config, + ) + def test_eval_step_calls_student_forward(self): """Verifies eval step correctly calls the student forward function and computes loss.""" # 1. Initialize Trainer (bypass init) diff --git a/tests/post_training/unit/train_sft_test.py b/tests/post_training/unit/train_sft_test.py index f0df897b3e..f579f8159d 100644 --- a/tests/post_training/unit/train_sft_test.py +++ b/tests/post_training/unit/train_sft_test.py @@ -31,6 +31,7 @@ class TrainSFTTest(unittest.TestCase): def test_validate_config_valid(self): config = SimpleNamespace( optimizer_memory_host_offload=False, + eval_interval=-1, ) # Should not raise any exception train_sft.validate_config(config) @@ -43,6 +44,119 @@ def test_validate_config_invalid_offload(self): with self.assertRaisesRegex(ValueError, "optimizer_memory_host_offload=True is not supported"): train_sft.validate_config(config) + @pytest.mark.cpu_only + def test_validate_config_accepts_weighted_diffusion_accumulation(self): + config = SimpleNamespace( + optimizer_memory_host_offload=False, + training_objective="block_diffusion", + gradient_accumulation_steps=2, + eval_interval=-1, + ) + + train_sft.validate_config(config) + + @pytest.mark.cpu_only + def test_validate_config_accepts_weighted_diffusion_evaluation(self): + config = SimpleNamespace( + optimizer_memory_host_offload=False, + training_objective="block_diffusion", + gradient_accumulation_steps=1, + eval_interval=10, + ) + + train_sft.validate_config(config) + + @pytest.mark.cpu_only + def test_setup_validates_before_initializing_model(self): + config = SimpleNamespace( + optimizer_memory_host_offload=True, + ) + + with ( + mock.patch.object(train_sft, "get_tunix_config") as get_tunix_config, + mock.patch.object(train_sft.model_creation_utils, "from_pretrained") as from_pretrained, + self.assertRaisesRegex(ValueError, "optimizer_memory_host_offload=True"), + ): + train_sft.setup_trainer_state(config) + + get_tunix_config.assert_not_called() + from_pretrained.assert_not_called() + + @pytest.mark.cpu_only + def test_maxtext_loss_wrapper_forwards_diffusion_masks(self): + trainer = mock.MagicMock() + trainer.with_loss_fn.return_value = trainer + config = SimpleNamespace() + train_sft.use_maxtext_loss_function(trainer, config) + wrapped_loss = trainer.with_loss_fn.call_args.args[0] + arrays = { + name: mock.sentinel.__getattr__(name) + for name in ( + "inputs", + "inputs_position", + "inputs_segmentation", + "targets", + "targets_position", + "targets_segmentation", + "completion_mask", + "corruption_mask", + "targets_loss_mask", + ) + } + + with mock.patch.object(train_sft, "loss_fn", return_value=(1.0, {})) as loss_fn_mock: + wrapped_loss(mock.sentinel.model, **arrays) + + forwarded = loss_fn_mock.call_args.args[2] + self.assertEqual(forwarded, arrays) + trainer.with_loss_fn.assert_called_once_with(wrapped_loss, has_aux=True) + + @pytest.mark.cpu_only + def test_diffusion_objective_uses_tunix_adapter(self): + trainer = mock.sentinel.trainer + config = SimpleNamespace(training_objective="block_diffusion") + with ( + mock.patch.object(train_sft.maxtext_diffusion_sft, "create_batch_adapter", return_value="adapter"), + mock.patch.object( + train_sft.maxtext_diffusion_sft, + "create_target_aligned_logits_fn", + return_value="logits_fn", + ), + mock.patch.object( + train_sft.tunix_diffusion_sft, + "configure_diffusion_sft", + return_value=trainer, + ) as configure, + ): + result = train_sft.configure_training_objective(trainer, config) + + self.assertIs(result, trainer) + configure.assert_called_once_with(trainer, "adapter", "logits_fn") + + @pytest.mark.cpu_only + def test_ar_objective_retains_maxtext_trainer_and_loss(self): + config = SimpleNamespace(training_objective="causal_lm") + with ( + mock.patch.object(train_sft, "MaxTextPeftTrainer", return_value=mock.sentinel.trainer) as trainer_type, + mock.patch.object(train_sft, "use_maxtext_loss_function", return_value=mock.sentinel.configured) as configure, + ): + trainer = train_sft._create_trainer("model", "optimizer", "config", config) + result = train_sft.configure_training_objective(trainer, config) + + self.assertIs(trainer, mock.sentinel.trainer) + self.assertIs(result, mock.sentinel.configured) + trainer_type.assert_called_once_with("model", "optimizer", "config") + configure.assert_called_once_with(mock.sentinel.trainer, config) + + @pytest.mark.cpu_only + def test_diffusion_objective_uses_current_tunix_trainer(self): + config = SimpleNamespace(training_objective="block_diffusion") + with mock.patch.object(train_sft.peft_trainer, "PeftTrainer", return_value=mock.sentinel.trainer) as trainer_type: + trainer = train_sft._create_trainer("model", "optimizer", "config", config) + + self.assertIs(trainer, mock.sentinel.trainer) + trainer_type.assert_called_once_with("model", "optimizer", "config") + @pytest.mark.cpu_only def test_train_model_caching_moe(self): """Test that NNX graph caching is disabled for MoE models (num_experts > 1).""" diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 20522bbd3b..0c7aa03005 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -44,9 +44,12 @@ from maxtext.layers import attention_op from maxtext.layers.attention_op import ( AttentionOp, + BlockCausalMask, ChunkedCausalMask, + _generate_block_causal_attention_mask, _generate_chunk_attention_mask, _make_bidirectional_block_mask, + _resolve_attention_type, ) from maxtext.layers.attentions import Attention from maxtext.layers import embeddings @@ -330,6 +333,143 @@ def test_value_error_on_zero_chunk_size(self): _generate_chunk_attention_mask(mask_shape=(4, 4), chunk_size=0) +class BlockCausalMaskTest(unittest.TestCase): + """Tests the shared dense and Splash block-causal masks.""" + + def _make_op(self, sequence_length, *, attention_type=AttentionType.BLOCK_DIFFUSION): + """Builds a minimal dot-product attention operator.""" + config = types.SimpleNamespace( + block_diffusion_block_size=4, + context_parallel_load_balance=False, + context_sharding="context", + ) + mesh = types.SimpleNamespace(shape={}) + kwargs = {} + if attention_type is not None: + kwargs["attention_type"] = attention_type + return AttentionOp( + config=config, + num_query_heads=1, + num_kv_heads=1, + max_target_length=sequence_length, + mesh=mesh, + attention_kernel="dot_product", + **kwargs, + ) + + def test_dense_and_splash_masks_match(self): + sequence_length = 10 + block_size = 4 + query_positions = np.arange(sequence_length)[:, None] + key_positions = np.arange(sequence_length)[None, :] + expected = query_positions // block_size >= key_positions // block_size + + splash_mask = BlockCausalMask((sequence_length, sequence_length), block_size) + dense_mask = _generate_block_causal_attention_mask((sequence_length, sequence_length), block_size) + + np.testing.assert_array_equal(splash_mask[:, :], expected) + np.testing.assert_array_equal(dense_mask, expected) + self.assertTrue(expected[1, 3]) + self.assertFalse(expected[3, 4]) + self.assertTrue(expected[4, 3]) + self.assertTrue(expected[8, 9]) + + def test_dot_product_mask_respects_packed_segments(self): + sequence_length = 12 + query = jnp.zeros((1, sequence_length, 1, 8)) + key = jnp.zeros((1, sequence_length, 1, 8)) + segment_ids = jnp.asarray([[1] * 8 + [2] * 4], dtype=jnp.int32) + + mask = self._make_op(sequence_length).generate_attention_mask( + query, + key, + segment_ids, + MODEL_MODE_TRAIN, + ) + + positions = np.arange(sequence_length) + expected = (positions[:, None] // 4 >= positions[None, :] // 4) & ( + np.asarray(segment_ids[0])[:, None] == np.asarray(segment_ids[0])[None, :] + ) + np.testing.assert_array_equal(np.asarray(mask == 0.0)[0, 0, 0], expected) + + def test_autoregressive_mask_is_unchanged(self): + key_length = 8 + query = jnp.zeros((1, 1, 1, 8)) + key = jnp.zeros((1, key_length, 1, 8)) + segment_ids = jnp.asarray([[1, 1, 0, 1, 0, 0, 1, 0]], dtype=jnp.int32) + + default_mask = self._make_op(key_length, attention_type=None).generate_attention_mask( + query, + key, + segment_ids, + MODEL_MODE_AUTOREGRESSIVE, + ) + block_diffusion_mask = self._make_op(key_length).generate_attention_mask( + query, + key, + segment_ids, + MODEL_MODE_AUTOREGRESSIVE, + ) + + np.testing.assert_array_equal(block_diffusion_mask, default_mask) + np.testing.assert_array_equal( + np.asarray(default_mask == 0.0)[0, 0, 0, 0], + np.asarray(segment_ids[0] == DECODING_ACTIVE_SEQUENCE_INDICATOR), + ) + + def test_invalid_block_size_raises(self): + with self.assertRaises(ValueError): + BlockCausalMask((4, 4), 0) + with self.assertRaises(ValueError): + _generate_block_causal_attention_mask((4, 4), -1) + + +class AttentionTypeResolutionTest(unittest.TestCase): + + def test_config_selects_block_diffusion_without_model_dispatch(self): + config = types.SimpleNamespace(attention_type=AttentionType.BLOCK_DIFFUSION.value) + + self.assertEqual(_resolve_attention_type(config, None), AttentionType.BLOCK_DIFFUSION) + self.assertEqual(_resolve_attention_type(config, AttentionType.GLOBAL), AttentionType.BLOCK_DIFFUSION) + self.assertEqual(_resolve_attention_type(config, AttentionType.FULL), AttentionType.FULL) + self.assertEqual(_resolve_attention_type(config, AttentionType.LOCAL_SLIDING), AttentionType.LOCAL_SLIDING) + self.assertEqual( + _resolve_attention_type(types.SimpleNamespace(attention_type=AttentionType.GLOBAL.value), AttentionType.FULL), + AttentionType.FULL, + ) + self.assertEqual(_resolve_attention_type(types.SimpleNamespace(), None), AttentionType.GLOBAL) + + def test_attention_op_honors_config_without_overriding_specialized_layers(self): + config = types.SimpleNamespace( + attention_type=AttentionType.BLOCK_DIFFUSION.value, + block_diffusion_block_size=4, + ) + op = AttentionOp( + config=config, + num_query_heads=1, + num_kv_heads=1, + max_target_length=8, + mesh=types.SimpleNamespace(shape={}), + attention_kernel="dot_product", + attention_type=AttentionType.GLOBAL, + ) + + self.assertEqual(op.attention_type, AttentionType.BLOCK_DIFFUSION) + + full_op = AttentionOp( + config=config, + num_query_heads=1, + num_kv_heads=1, + max_target_length=8, + mesh=types.SimpleNamespace(shape={}), + attention_kernel="dot_product", + attention_type=AttentionType.FULL, + ) + + self.assertEqual(full_op.attention_type, AttentionType.FULL) + + class LoadBalancedMaskTest(unittest.TestCase): """Tests for load-balanced Splash masks.""" @@ -370,6 +510,20 @@ def test_load_balanced_chunk_window(self): np.testing.assert_array_equal((causal_mask & chunk_mask)[:, :], expected_mask) + def test_load_balanced_block_causal_mask(self): + sequence_length = 8 + block_size = 2 + mask = attention_op.LoadBalancedBlockCausalMask( + shape=(sequence_length, sequence_length), + block_size=block_size, + cp_size=2, + ) + query_positions = mask.q_sequence[:, None] + key_positions = np.arange(sequence_length)[None, :] + expected = query_positions // block_size >= key_positions // block_size + + np.testing.assert_array_equal(mask[:, :], expected) + def test_dot_product_local_mask_uses_segment_positions(self): config = types.SimpleNamespace(context_parallel_load_balance=True, context_sharding="context") mesh = types.SimpleNamespace(shape={"context": 4}) diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index f0504e4dd8..a4e8abad40 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -192,6 +192,158 @@ 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_config(self): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "steps=1", + "attention=dot_product", + "attention_type=block_diffusion", + "block_diffusion_block_size=7", + "vocab_size=256", + "max_target_length=2048", + "hardware=cpu", + ] + config = pyconfig.initialize(argv) + + self.assertEqual(config.attention_type, "block_diffusion") + self.assertEqual(config.block_diffusion_block_size, 7) + self.assertNotEqual(config.max_target_length % config.block_diffusion_block_size, 0) + self.assertFalse(hasattr(config, "enable_block_diffusion")) + + def test_block_diffusion_invalid_configs_raise(self): + base_overrides = { + "run_name": "test", + "steps": 1, + "attention": "dot_product", + "attention_type": "block_diffusion", + "block_diffusion_block_size": 32, + "vocab_size": 256, + "max_target_length": 2048, + "hardware": "cpu", + } + cases = [ + {"block_diffusion_block_size": 0}, + {"attention": "cudnn_flash_te"}, + {"attention": "flash", "hardware": "gpu"}, + ] + for overrides in cases: + with self.subTest(overrides=overrides): + config = base_overrides | overrides + argv = ["", _BASE_CONFIG_PATH, *(f"{key}={value}" for key, value in config.items())] + with self.assertRaises((ValueError, pydantic.ValidationError)): + pyconfig.initialize(argv) + + def test_default_attention_remains_global(self): + config = pyconfig.initialize(["", _BASE_CONFIG_PATH, "run_name=test", "steps=1"]) + + self.assertEqual(config.attention_type, "global") + self.assertEqual(config.training_objective, "causal_lm") + self.assertEqual(config.block_diffusion_mask_id, -1) + self.assertEqual(config.distill_data_source, "dataset") + self.assertEqual(config.distill_rollout_algorithm, "low_confidence") + self.assertEqual(config.distill_rollout_stop_token_ids, []) + + def test_block_diffusion_training_objective_config(self): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "steps=1", + "training_objective=block_diffusion", + "attention=dot_product", + "attention_type=block_diffusion", + "block_diffusion_block_size=7", + "block_diffusion_mask_id=100", + "block_diffusion_min_noise=0.05", + "vocab_size=256", + "max_target_length=2048", + "packing=False", + "dataset_type=hf", + "hf_path=parquet", + "hardware=cpu", + ] + + config = pyconfig.initialize(argv) + + self.assertEqual(config.training_objective, "block_diffusion") + self.assertEqual(config.block_diffusion_mask_id, 100) + self.assertEqual(config.block_diffusion_min_noise, 0.05) + self.assertEqual(config.block_diffusion_logit_alignment, "same_position") + self.assertEqual(config.block_diffusion_canvas_policy, "all_masked") + self.assertNotEqual(config.max_target_length % config.block_diffusion_block_size, 0) + + def test_block_diffusion_training_objective_rejects_incompatible_config(self): + base_overrides = { + "run_name": "test", + "steps": 1, + "training_objective": "block_diffusion", + "attention": "dot_product", + "attention_type": "block_diffusion", + "block_diffusion_block_size": 32, + "block_diffusion_mask_id": 100, + "block_diffusion_min_noise": 0.05, + "vocab_size": 256, + "max_target_length": 2048, + "packing": False, + "dataset_type": "hf", + "hf_path": "parquet", + "hardware": "cpu", + } + cases = [ + ({"attention_type": "global"}, "attention_type='block_diffusion'"), + ({"block_diffusion_mask_id": -1}, "block_diffusion_mask_id"), + ({"block_diffusion_mask_id": 256}, "block_diffusion_mask_id"), + ({"block_diffusion_min_noise": 0.0}, "block_diffusion_min_noise"), + ({"packing": True}, "packing=False"), + ({"mtp_num_layers": 1}, "MTP"), + ({"num_vocab_tiling": 2}, "vocabulary tiling"), + ({"dataset_type": "grain"}, "dataset_type='hf'"), + ({"use_dpo": True}, "DPO"), + ({"use_multimodal": True}, "text-only"), + ({"block_diffusion_logit_alignment": "shifted"}, "seed_and_mask"), + ( + {"block_diffusion_canvas_policy": "seed_and_mask"}, + "same_position/all_masked", + ), + ] + for overrides, expected_regex in cases: + with self.subTest(overrides=overrides): + config = base_overrides | overrides + argv = [ + "", + _BASE_CONFIG_PATH, + *(f"{key}={value}" for key, value in config.items()), + ] + with self.assertRaisesRegex((ValueError, pydantic.ValidationError), expected_regex): + pyconfig.initialize(argv) + + def test_shifted_block_diffusion_requires_seeded_canvas(self): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "steps=1", + "training_objective=block_diffusion", + "attention=dot_product", + "attention_type=block_diffusion", + "block_diffusion_block_size=8", + "block_diffusion_mask_id=100", + "block_diffusion_logit_alignment=shifted", + "block_diffusion_canvas_policy=seed_and_mask", + "vocab_size=256", + "packing=False", + "dataset_type=hf", + "hf_path=parquet", + "hardware=cpu", + ] + + config = pyconfig.initialize(argv) + + self.assertEqual(config.block_diffusion_logit_alignment, "shifted") + self.assertEqual(config.block_diffusion_canvas_policy, "seed_and_mask") + @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/diffusion_denoise_test.py b/tests/unit/diffusion_denoise_test.py new file mode 100644 index 0000000000..d00df168e5 --- /dev/null +++ b/tests/unit/diffusion_denoise_test.py @@ -0,0 +1,456 @@ +# Copyright 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 +# +# http://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. + +"""Tests model-independent block-diffusion rollout transitions.""" + +from absl.testing import absltest +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.diffusion import denoise + + +def _target_logits(targets, vocab_size=32, high=12.0): + return jax.nn.one_hot(targets, vocab_size, dtype=jnp.float32) * high + + +class DiffusionDenoiseTest(absltest.TestCase): + + def test_trace_records_each_pre_commit_step_and_logprob(self): + initial = jnp.asarray([[4, 1, 1, 1]], dtype=jnp.int32) + positions = jnp.arange(4, dtype=jnp.int32)[None, :] + completion = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_) + + trace = denoise.low_confidence_rollout( + lambda canvas: jnp.zeros((*canvas.shape, 8), dtype=jnp.float32), + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=7, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + ) + + np.testing.assert_array_equal(trace.action_steps, [[-1, 0, 1, 2]]) + np.testing.assert_allclose(trace.action_logps[:, 1:], -np.log(7.0), rtol=1e-6) + self.assertFalse(bool(jnp.any(trace.tokens == 7))) + + def test_trace_records_shifted_anchor_as_an_action(self): + initial = jnp.asarray([[2, 3, 4, 5, 1, 1]], dtype=jnp.int32) + positions = jnp.arange(6, dtype=jnp.int32)[None, :] + completion = jnp.asarray([[0, 0, 0, 0, 1, 1]], dtype=jnp.bool_) + + trace = denoise.low_confidence_rollout( + lambda canvas: _target_logits(jnp.full_like(canvas, 6), vocab_size=8), + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=7, + logit_alignment="shifted", + canvas_policy="seed_and_mask", + ) + + np.testing.assert_array_equal(trace.action_steps, [[-1, -1, -1, -1, 0, 1]]) + np.testing.assert_array_equal(trace.tokens, [[2, 3, 4, 5, 6, 6]]) + + def test_sampled_trace_is_seeded_and_excludes_mask_token(self): + initial = jnp.asarray([[4, 1, 1, 1, 1, 1, 1, 1]], dtype=jnp.int32) + positions = jnp.arange(8, dtype=jnp.int32)[None, :] + completion = jnp.asarray([[0, 1, 1, 1, 1, 1, 1, 1]], dtype=jnp.bool_) + + def rollout(seed): + return denoise.low_confidence_rollout( + lambda canvas: jnp.zeros((*canvas.shape, 32), dtype=jnp.float32), + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=31, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + rng=jax.random.PRNGKey(seed), + ) + + first = rollout(7) + repeated = rollout(7) + different = rollout(8) + + np.testing.assert_array_equal(first.tokens, repeated.tokens) + self.assertFalse(bool(jnp.array_equal(first.tokens, different.tokens))) + self.assertFalse(bool(jnp.any(first.tokens == 31))) + + def test_rollout_accepts_typed_and_legacy_keys_for_each_prng_implementation(self): + initial = jnp.asarray([[4, 7, 7, 7], [3, 7, 7, 7]], dtype=jnp.int32) + positions = jnp.broadcast_to(jnp.arange(4, dtype=jnp.int32), initial.shape) + completion = positions > 0 + + def logits_fn(canvas): + token_ids = (jnp.arange(canvas.shape[1])[None, :] + 1) % 7 + return jax.nn.one_hot(jnp.broadcast_to(token_ids, canvas.shape), 8, dtype=jnp.float32) * 4.0 + + for implementation in ("threefry2x32", "unsafe_rbg"): + with ( + self.subTest(implementation=implementation), + jax.default_prng_impl(implementation), + ): + legacy_key = jax.random.PRNGKey(7) + typed_key = jax.random.key(7) + keys = { + "legacy_scalar": legacy_key, + "legacy_batched": jax.vmap(lambda row: jax.random.fold_in(legacy_key, row))(jnp.arange(2)), + "typed_scalar": typed_key, + "typed_batched": jax.vmap(lambda row: jax.random.fold_in(typed_key, row))(jnp.arange(2)), + } + traces = {} + for key_kind, key in keys.items(): + with self.subTest(implementation=implementation, key_kind=key_kind): + traces[key_kind] = denoise.low_confidence_rollout( + logits_fn, + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=7, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + rng=key, + ) + self.assertFalse(bool(jnp.any(traces[key_kind].tokens == 7))) + + for trace in traces.values(): + np.testing.assert_array_equal(trace.tokens, traces["legacy_scalar"].tokens) + np.testing.assert_array_equal(trace.action_steps, traces["legacy_scalar"].action_steps) + + generated = denoise.low_confidence_generate( + logits_fn, + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=7, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + ) + self.assertFalse(bool(jnp.any(generated == 7))) + + def test_trace_steps_are_compact_per_row_with_heterogeneous_prompts(self): + initial = jnp.asarray( + [ + [0, 0, 0, 4, 1, 1, 1, 1], + [4, 5, 6, 7, 1, 1, 1, 1], + ], + dtype=jnp.int32, + ) + validity = initial != 0 + positions = jnp.where(validity, jnp.cumsum(validity, axis=1) - 1, 0) + completion = jnp.asarray( + [ + [0, 0, 0, 0, 1, 1, 1, 1], + [0, 0, 0, 0, 1, 1, 1, 1], + ], + dtype=jnp.bool_, + ) + + trace = denoise.low_confidence_rollout( + lambda canvas: jnp.zeros((*canvas.shape, 16), dtype=jnp.float32), + initial, + positions, + validity, + completion, + block_size=4, + mask_id=15, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + ) + + np.testing.assert_array_equal(trace.action_steps[:, 4:], [[0, 1, 2, 3], [0, 1, 2, 3]]) + + def test_row_sampling_is_independent_of_other_rows_denoise_progress(self): + first = jnp.asarray([[0, 0, 0, 4, 1, 1, 1, 1]], dtype=jnp.int32) + second = jnp.asarray([[4, 5, 6, 7, 1, 1, 1, 1]], dtype=jnp.int32) + + def rollout(initial): + validity = initial != 0 + positions = jnp.where(validity, jnp.cumsum(validity, axis=1) - 1, 0) + completion = jnp.zeros_like(validity).at[:, 4:].set(True) + return denoise.low_confidence_rollout( + lambda canvas: jnp.zeros((*canvas.shape, 16), dtype=jnp.float32), + initial, + positions, + validity, + completion, + block_size=4, + mask_id=15, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + rng=jax.random.PRNGKey(11), + ) + + alone = rollout(first) + batched = rollout(jnp.concatenate([first, second], axis=0)) + + np.testing.assert_array_equal(alone.tokens[0], batched.tokens[0]) + np.testing.assert_allclose(alone.action_logps[0], batched.action_logps[0]) + + def test_same_position_generates_partial_blocks_and_preserves_prompt(self): + initial = jnp.asarray([[7, 8, 1, 1, 1, 1, 1, 0]], dtype=jnp.int32) + expected = jnp.asarray([[7, 8, 12, 13, 14, 15, 16, 0]], dtype=jnp.int32) + positions = jnp.arange(8, dtype=jnp.int32)[None, :] + validity = initial != 0 + completion = jnp.asarray([[0, 0, 1, 1, 1, 1, 1, 0]], dtype=jnp.bool_) + + generated = denoise.low_confidence_generate( + lambda _: _target_logits(expected), + initial, + positions, + validity, + completion, + block_size=4, + mask_id=31, + logit_alignment="same_position", + canvas_policy="all_masked", + ) + + np.testing.assert_array_equal(generated, expected) + + def test_shifted_seed_generates_anchors_before_each_block(self): + initial = jnp.asarray([[5, 6, 1, 1, 1, 1, 1, 0]], dtype=jnp.int32) + expected = jnp.asarray([[5, 6, 9, 10, 11, 12, 13, 0]], dtype=jnp.int32) + positions = jnp.arange(8, dtype=jnp.int32)[None, :] + validity = initial != 0 + completion = jnp.asarray([[0, 0, 1, 1, 1, 1, 1, 0]], dtype=jnp.bool_) + + generated = denoise.low_confidence_generate( + lambda _: _target_logits(expected), + initial, + positions, + validity, + completion, + block_size=4, + mask_id=31, + logit_alignment="shifted", + canvas_policy="seed_and_mask", + ) + + np.testing.assert_array_equal(generated, expected) + self.assertEqual(int(generated[0, 4]), 11) + + def test_forced_argmax_progress_completes_low_confidence_rows(self): + initial = jnp.asarray([[4, 1, 1, 1]], dtype=jnp.int32) + positions = jnp.arange(4, dtype=jnp.int32)[None, :] + completion = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_) + + generated = denoise.low_confidence_generate( + lambda canvas: jnp.zeros((*canvas.shape, 8), dtype=jnp.float32), + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=7, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + ) + + self.assertFalse(bool(jnp.any(generated == 7))) + np.testing.assert_array_equal(generated[:, :1], initial[:, :1]) + + def test_mask_token_is_never_committed(self): + initial = jnp.asarray([[6, 1, 1, 1]], dtype=jnp.int32) + positions = jnp.arange(4, dtype=jnp.int32)[None, :] + completion = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_) + + def mask_favoring_logits(canvas): + logits = jnp.zeros((*canvas.shape, 8), dtype=jnp.float32) + return logits.at[..., 7].set(100.0).at[..., 5].set(10.0) + + generated = denoise.low_confidence_generate( + mask_favoring_logits, + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=7, + logit_alignment="same_position", + canvas_policy="all_masked", + ) + + np.testing.assert_array_equal(generated, [[6, 5, 5, 5]]) + + def test_mask_token_is_excluded_from_shifted_anchor_proposals(self): + initial = jnp.asarray([[2, 3, 4, 5, 1, 1]], dtype=jnp.int32) + positions = jnp.arange(6, dtype=jnp.int32)[None, :] + completion = jnp.asarray([[0, 0, 0, 0, 1, 1]], dtype=jnp.bool_) + + def mask_favoring_logits(canvas): + logits = jnp.zeros((*canvas.shape, 8), dtype=jnp.float32) + return logits.at[..., 7].set(100.0).at[..., 6].set(10.0) + + generated = denoise.low_confidence_generate( + mask_favoring_logits, + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=7, + logit_alignment="shifted", + canvas_policy="seed_and_mask", + ) + + np.testing.assert_array_equal(generated, [[2, 3, 4, 5, 6, 6]]) + + def test_rollout_is_jittable(self): + positions = jnp.arange(4, dtype=jnp.int32)[None, :] + completion = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_) + expected = jnp.asarray([[4, 5, 6, 7]], dtype=jnp.int32) + + generate = jax.jit( + lambda initial: denoise.low_confidence_generate( + lambda _: _target_logits(expected), + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=31, + logit_alignment="same_position", + canvas_policy="all_masked", + ) + ) + + np.testing.assert_array_equal(generate(jnp.asarray([[4, 1, 1, 1]], dtype=jnp.int32)), expected) + + def test_logical_positions_drive_blocks_after_reordering(self): + positions = jnp.asarray([[0, 4, 1, 5, 2, 6, 3, 7]], dtype=jnp.int32) + expected = positions + 10 + initial = jnp.where(positions < 2, expected, 1) + completion = positions >= 2 + + generated = denoise.low_confidence_generate( + lambda _: _target_logits(expected), + initial, + positions, + jnp.ones_like(completion), + completion, + block_size=4, + mask_id=31, + logit_alignment="same_position", + canvas_policy="all_masked", + ) + + np.testing.assert_array_equal(generated, expected) + + def test_forced_ties_are_invariant_to_physical_reordering(self): + completion = jnp.asarray([[0, 1, 1, 1]], dtype=jnp.bool_) + + def generate(physical_order): + positions = jnp.asarray(physical_order, dtype=jnp.int32)[None, :] + initial = jnp.take_along_axis(jnp.asarray([[4, 1, 1, 1]], dtype=jnp.int32), positions, axis=1) + physical_completion = jnp.take_along_axis(completion, positions, axis=1) + + def low_confidence_logits(canvas): + committed_count = jnp.sum((canvas != 31) & physical_completion, axis=1) + proposed_token = jnp.broadcast_to((10 + committed_count)[:, None], canvas.shape) + return _target_logits(proposed_token, vocab_size=32, high=0.1) + + output = denoise.low_confidence_generate( + low_confidence_logits, + initial, + positions, + jnp.ones_like(physical_completion), + physical_completion, + block_size=4, + mask_id=31, + logit_alignment="same_position", + canvas_policy="all_masked", + confidence_threshold=0.99, + ) + inverse_order = jnp.argsort(positions, axis=1) + return jnp.take_along_axis(output, inverse_order, axis=1) + + np.testing.assert_array_equal(generate([0, 1, 2, 3]), generate([0, 3, 1, 2])) + + def test_rejects_non_suffix_and_shifted_origin_completion(self): + positions = np.arange(4, dtype=np.int32)[None, :] + validity = np.ones((1, 4), dtype=bool) + + with self.assertRaisesRegex(ValueError, "contiguous suffix"): + denoise.validate_completion_suffix(positions, validity, np.asarray([[0, 1, 0, 1]], dtype=bool)) + with self.assertRaisesRegex(ValueError, "position zero"): + denoise.validate_completion_suffix( + positions, + validity, + np.asarray([[1, 1, 1, 1]], dtype=bool), + shifted_seed=True, + ) + + def test_rejects_missing_duplicate_or_out_of_range_positions(self): + validity = np.ones((1, 3), dtype=bool) + completion = np.asarray([[0, 1, 1]], dtype=bool) + + for positions in ( + np.asarray([[1, 2, 3]], dtype=np.int32), + np.asarray([[0, 1, 1]], dtype=np.int32), + np.asarray([[0, 1, 5]], dtype=np.int32), + ): + with self.subTest(positions=positions), self.assertRaisesRegex(ValueError, "logical positions"): + denoise.validate_completion_suffix(positions, validity, completion) + + def test_rejects_unsupported_contract_and_short_step_cap(self): + tokens = jnp.ones((1, 4), dtype=jnp.int32) + positions = jnp.arange(4, dtype=jnp.int32)[None, :] + mask = jnp.ones_like(tokens, dtype=jnp.bool_) + + with self.assertRaisesRegex(ValueError, "supports only"): + denoise.low_confidence_generate( + lambda _: _target_logits(tokens), + tokens, + positions, + mask, + mask, + block_size=4, + mask_id=31, + logit_alignment="same_position", + canvas_policy="seed_and_mask", + ) + with self.assertRaisesRegex(ValueError, "at least block_size"): + denoise.low_confidence_generate( + lambda _: _target_logits(tokens), + tokens, + positions, + mask, + mask, + block_size=4, + mask_id=31, + logit_alignment="same_position", + canvas_policy="all_masked", + max_denoise_steps=3, + ) diff --git a/tests/unit/gradient_accumulation_nnx_test.py b/tests/unit/gradient_accumulation_nnx_test.py index 478fe9a746..5ca56f484d 100644 --- a/tests/unit/gradient_accumulation_nnx_test.py +++ b/tests/unit/gradient_accumulation_nnx_test.py @@ -14,6 +14,8 @@ """Unit tests for the NNX branch of gradient_accumulation_loss_and_grad.""" +# pylint: disable=too-many-positional-arguments + import unittest from dataclasses import dataclass @@ -67,6 +69,20 @@ def _fake_loss_fn(model, config, data, dropout_rng, params, is_train=True): return xent_sum / total_weights, aux +def _zero_weight_loss_fn(model, config, data, dropout_rng, params, is_train=True): + """Produces model-dependent sums with a zero token denominator.""" + del config, dropout_rng, params, is_train + xent_sum = jnp.sum(model(data["inputs"]) ** 2) + 1.0 + aux = { + "xent_sum": xent_sum, + "total_weights": jnp.array(0.0), + "moe_lb_loss": jnp.array(0.0), + "indexer_loss": jnp.array(0.0), + "mtp_loss": jnp.array(0.0), + } + return xent_sum, aux + + class TestGradientAccumulationNNX(unittest.TestCase): """Cover the NNX path of gradient_accumulation_loss_and_grad.""" @@ -151,6 +167,22 @@ def test_nnx_with_shard_optimizer_over_data_casts_to_bf16(self): for g in jax.tree.leaves(raw_grads): self.assertTrue(jnp.all(jnp.isfinite(g))) + def test_zero_total_weights_returns_zero_loss_and_gradients(self): + loss, _, raw_grads = gradient_accumulation.gradient_accumulation_loss_and_grad( + _zero_weight_loss_fn, + self.cfg, + self.model, + params=None, + params_shardings=self._params_shardings(), + data=self.data, + dropout_rng=None, + ) + + self.assertEqual(float(loss), 0.0) + for gradient in jax.tree.leaves(raw_grads): + self.assertTrue(jnp.all(jnp.isfinite(gradient))) + self.assertTrue(jnp.all(gradient == 0)) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/hf_data_processing_test.py b/tests/unit/hf_data_processing_test.py index 93564514a3..a23525af02 100644 --- a/tests/unit/hf_data_processing_test.py +++ b/tests/unit/hf_data_processing_test.py @@ -15,6 +15,7 @@ """Tests for Hugging Face data processing.""" import sys +from types import SimpleNamespace import unittest import os.path @@ -26,6 +27,7 @@ from maxtext.configs import pyconfig from maxtext.input_pipeline import hf_data_processing from maxtext.input_pipeline import input_pipeline_interface +from maxtext.input_pipeline import input_pipeline_utils from maxtext.common.gcloud_stub import is_decoupled from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory @@ -126,5 +128,75 @@ def get_first_batch(iterator): self.assertTrue((train_batch1["targets"] == train_batch2["targets"]).all()) # pytype: disable=unsupported-operands +@pytest.mark.cpu_only +class TrainingObjectiveTransformTest(unittest.TestCase): + """Tests objective selection at the HF pipeline boundary.""" + + def test_default_objective_keeps_shift_data(self): + transform = hf_data_processing._get_training_objective_transform( # pylint: disable=protected-access + SimpleNamespace(), + shift=True, + use_dpo=False, + use_sft=False, + completion_only=False, + packing=False, + pad_id=0, + bos_token_id=1, + ) + + self.assertIsInstance(transform, input_pipeline_utils.ShiftData) + self.assertEqual(transform.ignored_ids, [0, 1]) + + def test_explicit_block_diffusion_objective_replaces_shift(self): + config = SimpleNamespace( + training_objective="block_diffusion", + block_diffusion_block_size=4, + block_diffusion_mask_id=99, + block_diffusion_min_noise=0.05, + block_diffusion_canvas_policy="seed_and_mask", + block_diffusion_logit_alignment="shifted", + ) + + transform = hf_data_processing._get_training_objective_transform( # pylint: disable=protected-access + config, + shift=True, + use_dpo=False, + use_sft=True, + completion_only=True, + packing=False, + pad_id=0, + bos_token_id=1, + ) + + self.assertIsInstance(transform, input_pipeline_utils.BlockDiffusionCorruption) + self.assertEqual(transform.block_size, 4) + self.assertEqual(transform.mask_id, 99) + self.assertEqual(transform.min_noise, 0.05) + self.assertTrue(transform.completion_only) + self.assertTrue(transform.seed_first_token) + self.assertTrue(transform.include_seed_in_loss) + + def test_block_diffusion_rejects_packing(self): + config = SimpleNamespace( + training_objective="block_diffusion", + block_diffusion_block_size=4, + block_diffusion_mask_id=99, + block_diffusion_min_noise=0.05, + block_diffusion_canvas_policy="all_masked", + ) + + with self.assertRaisesRegex(ValueError, "packing=False"): + hf_data_processing._get_training_objective_transform( # pylint: disable=protected-access + config, + shift=True, + use_dpo=False, + use_sft=False, + completion_only=False, + packing=True, + pad_id=0, + bos_token_id=1, + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/input_pipeline_utils_test.py b/tests/unit/input_pipeline_utils_test.py index 0d5adae4a9..ec775a949e 100644 --- a/tests/unit/input_pipeline_utils_test.py +++ b/tests/unit/input_pipeline_utils_test.py @@ -17,7 +17,14 @@ 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 ( + BlockDiffusionCorruption, + PadOrTrimToMaxLength, + SFTPromptMasking, + compute_file_sharding, +) @pytest.mark.cpu_only @@ -109,5 +116,281 @@ def test_no_row_shard_when_only_one_reader(self): self.assertIsNone(row_shard) +def _make_block_diffusion_batch(): + """Builds a padded, completion-only SFT batch with a partial final block.""" + inputs = np.asarray( + [ + [11, 12, 21, 22, 23, 24, 25, 26, 27, 0], + [31, 41, 42, 43, 44, 45, 0, 0, 0, 0], + ], + dtype=np.int32, + ) + targets = np.asarray( + [ + [11, 12, 21, 22, 23, 24, 25, 26, 27, 0], + [31, 41, 42, 43, 44, 45, 0, 0, 0, 0], + ], + dtype=np.int32, + ) + completion_mask = np.asarray( + [ + [0, 0, 1, 1, 1, 1, 1, 1, 1, 0], + [0, 1, 1, 1, 1, 1, 0, 0, 0, 0], + ], + dtype=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 != 0).astype(np.int32), + "targets_segmentation": (targets != 0).astype(np.int32), + "inputs_position": positions.copy(), + "targets_position": positions.copy(), + "completion_mask": completion_mask, + } + + +class _ControlledRng: + """Supplies distinct noise rates while making token draws deterministic.""" + + def __init__(self): + self.uniform_calls = 0 + self.random_calls = 0 + + def uniform(self, low, high, size): + del low, high + self.uniform_calls += 1 + return np.asarray([[[0.1], [0.9], [0.1]], [[0.1], [0.9], [0.1]]]).reshape(size) + + def random(self, size): + self.random_calls += 1 + return np.full(size, 0.5) + + +@pytest.mark.cpu_only +class BlockDiffusionCorruptionTest(unittest.TestCase): + """Tests aligned block-diffusion corruption metadata.""" + + def _apply(self, seed=0): + transform = BlockDiffusionCorruption(block_size=4, mask_id=99, min_noise=1.0e-3, completion_only=True) + return transform.random_map(_make_block_diffusion_batch(), np.random.default_rng(seed)) + + def test_preserves_targets_and_segmentation(self): + clean = _make_block_diffusion_batch() + output = self._apply() + + for key in ( + "targets", + "inputs_segmentation", + "targets_segmentation", + "inputs_position", + "targets_position", + ): + np.testing.assert_array_equal(output[key], clean[key]) + np.testing.assert_array_equal(output["completion_mask"], clean["completion_mask"]) + self.assertTrue(np.all(output["targets_loss_mask"] <= output["completion_mask"])) + self.assertFalse(np.shares_memory(output["completion_mask"], output["targets_loss_mask"])) + np.testing.assert_array_equal(output["corruption_mask"], output["targets_loss_mask"]) + + loss_positions = output["targets_loss_mask"] != 0 + self.assertTrue(np.all(output["inputs"][loss_positions] == 99)) + np.testing.assert_array_equal(output["inputs"][~loss_positions], clean["inputs"][~loss_positions]) + + def test_every_eligible_full_or_partial_block_has_loss(self): + output = self._apply(seed=17) + completion_mask = output["completion_mask"] != 0 + loss_mask = output["targets_loss_mask"] != 0 + + for row in range(completion_mask.shape[0]): + for start in range(0, completion_mask.shape[1], 4): + stop = min(start + 4, completion_mask.shape[1]) + if np.any(completion_mask[row, start:stop]): + self.assertTrue(np.any(loss_mask[row, start:stop]), (row, start, stop)) + self.assertTrue(loss_mask[0, 8]) + + def test_corruption_rate_is_sampled_per_block(self): + rng = _ControlledRng() + transform = BlockDiffusionCorruption(block_size=4, mask_id=99, completion_only=True) + + output = transform.random_map(_make_block_diffusion_batch(), rng) + + self.assertEqual(rng.uniform_calls, 1) + self.assertEqual(rng.random_calls, 2) + self.assertEqual(int(output["targets_loss_mask"][0, :4].sum()), 1) + self.assertEqual(int(output["targets_loss_mask"][0, 4:8].sum()), 4) + + def test_deterministic_for_fixed_seed(self): + output_a = self._apply(seed=123) + output_b = self._apply(seed=123) + for key, value in output_a.items(): + np.testing.assert_array_equal(value, output_b[key]) + + def test_invalid_parameters_raise(self): + with self.assertRaises(ValueError): + BlockDiffusionCorruption(block_size=0, mask_id=99) + with self.assertRaises(ValueError): + BlockDiffusionCorruption(block_size=4, mask_id=99, min_noise=0.0) + with self.assertRaisesRegex(ValueError, "requires seed_first_token"): + BlockDiffusionCorruption(block_size=4, mask_id=99, include_seed_in_loss=True) + + def test_completion_only_requires_explicit_role_mask(self): + clean = _make_block_diffusion_batch() + del clean["completion_mask"] + + with self.assertRaisesRegex(ValueError, "explicit completion_mask"): + self._apply_transform(clean) + + def test_opd_suffix_contract_selects_final_assistant_turn(self): + clean = _make_block_diffusion_batch() + clean["completion_mask"][0] = np.asarray([0, 0, 1, 1, 0, 0, 1, 1, 1, 0]) + transform = BlockDiffusionCorruption( + block_size=4, + mask_id=99, + completion_only=True, + select_last_completion_suffix=True, + ) + + output = transform.random_map(clean, np.random.default_rng(0)) + + np.testing.assert_array_equal( + output["completion_mask"][0], + np.asarray([0, 0, 0, 0, 0, 0, 1, 1, 1, 0]), + ) + + def test_opd_suffix_contract_truncates_context_after_last_completion(self): + clean = _make_block_diffusion_batch() + clean["completion_mask"][0] = np.asarray([0, 0, 1, 1, 0, 0, 0, 0, 0, 0]) + transform = BlockDiffusionCorruption( + block_size=4, + mask_id=99, + completion_only=True, + select_last_completion_suffix=True, + ) + + output = transform.random_map(clean, np.random.default_rng(0)) + + np.testing.assert_array_equal( + output["completion_mask"][0], + np.asarray([0, 0, 1, 1, 0, 0, 0, 0, 0, 0]), + ) + np.testing.assert_array_equal( + output["inputs_segmentation"][0], + np.asarray([1, 1, 1, 1, 0, 0, 0, 0, 0, 0]), + ) + np.testing.assert_array_equal( + output["targets_segmentation"][0], + np.asarray([1, 1, 1, 1, 0, 0, 0, 0, 0, 0]), + ) + np.testing.assert_array_equal(output["targets"][0, 4:], 0) + + def test_opd_suffix_contract_requires_prompt_and_completion(self): + transform = BlockDiffusionCorruption( + block_size=4, + mask_id=99, + completion_only=True, + select_last_completion_suffix=True, + ) + for completion_mask in ( + np.zeros(10, dtype=np.int32), + np.asarray([1, 1, 1, 1, 1, 1, 1, 1, 1, 0]), + ): + clean = _make_block_diffusion_batch() + clean["completion_mask"][0] = completion_mask + with self.assertRaisesRegex(ValueError, "completion span after a prompt"): + transform.random_map(clean, np.random.default_rng(0)) + + def test_opd_suffix_contract_accepts_prompt_then_completion(self): + transform = BlockDiffusionCorruption( + block_size=4, + mask_id=99, + completion_only=True, + select_last_completion_suffix=True, + ) + + output = transform.random_map(_make_block_diffusion_batch(), np.random.default_rng(0)) + + self.assertTrue(output["targets_loss_mask"].any()) + + def test_completion_mask_does_not_replace_full_validity(self): + clean = _make_block_diffusion_batch() + transform = BlockDiffusionCorruption(block_size=4, mask_id=99, min_noise=1.0, completion_only=False) + + output = transform.random_map(clean, np.random.default_rng(0)) + + np.testing.assert_array_equal(output["targets_segmentation"], clean["targets_segmentation"]) + np.testing.assert_array_equal(output["completion_mask"], clean["completion_mask"]) + self.assertTrue(np.all(output["targets_loss_mask"][clean["inputs_segmentation"] != 0])) + self.assertTrue(np.any(output["targets_loss_mask"] > output["completion_mask"])) + + def test_all_padding_has_finite_zero_masks(self): + clean = _make_block_diffusion_batch() + for key in ( + "inputs", + "targets", + "inputs_segmentation", + "targets_segmentation", + "completion_mask", + ): + clean[key] = np.zeros_like(clean[key]) + + output = self._apply_transform(clean) + + self.assertFalse(output["corruption_mask"].any()) + self.assertFalse(output["targets_loss_mask"].any()) + + def test_shifted_seed_canvas_keeps_anchors_clean_but_supervised(self): + clean = _make_block_diffusion_batch() + transform = BlockDiffusionCorruption( + block_size=4, + mask_id=99, + min_noise=1.0, + completion_only=False, + seed_first_token=True, + include_seed_in_loss=True, + ) + + output = transform.random_map(clean, np.random.default_rng(0)) + + self.assertFalse(output["corruption_mask"][:, ::4].any()) + expected_loss = clean["targets_segmentation"].copy() + expected_loss[:, 0] = 0 + np.testing.assert_array_equal(output["targets_loss_mask"], expected_loss) + self.assertTrue(output["targets_loss_mask"][0, 4]) + self.assertTrue(output["targets_loss_mask"][0, 8]) + + def _apply_transform(self, batch): + transform = BlockDiffusionCorruption(block_size=4, mask_id=99, completion_only=True) + return transform.random_map(batch, np.random.default_rng(0)) + + +@pytest.mark.cpu_only +class BlockDiffusionSftMaskTest(unittest.TestCase): + """Tests role metadata before randomized corruption.""" + + def test_target_aligned_sft_preserves_clean_tokens_and_completion_roles(self): + transform = SFTPromptMasking( + text_column_name="messages", + completion_only=True, + max_target_length=8, + unk_id=0, + target_aligned=True, + emit_completion_mask=True, + ) + element = { + "messages": [[11, 12], [21, 22], [31], [41, 42]], + "is_prompt": [True, False, True, False], + } + + output = transform.map(element) + padded = PadOrTrimToMaxLength(max_length=8, pad_id=0).map(output) + + np.testing.assert_array_equal(padded["targets"], padded["inputs"]) + np.testing.assert_array_equal(padded["completion_mask"], [0, 0, 1, 1, 0, 1, 1, 0]) + np.testing.assert_array_equal(padded["targets_segmentation"], [1, 1, 1, 1, 1, 1, 1, 0]) + self.assertNotIn("completion_mask_segmentation", padded) + self.assertNotIn("completion_mask_position", padded) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/max_utils_test.py b/tests/unit/max_utils_test.py index 097d85d491..7dcc6cea83 100644 --- a/tests/unit/max_utils_test.py +++ b/tests/unit/max_utils_test.py @@ -23,10 +23,12 @@ import jax from jax import numpy as jnp from jax import random +from maxtext.common.common_types import ReorderStrategy from maxtext.configs import pyconfig from maxtext.utils import max_utils from maxtext.utils.train_utils import setup_train_loop from tests.utils.test_helpers import get_test_config_path +import numpy as np import optax import pytest @@ -655,6 +657,31 @@ def test_reorder_roundtrip(self): # 3. Assert roundtrip is lossless self.assertTrue(jnp.allclose(x, restored, rtol=1e-5, atol=1e-6)) + def test_block_diffusion_masks_are_reordered_with_tokens(self): + completion_mask = jnp.arange(16, dtype=jnp.int32).reshape(1, 16) + targets_loss_mask = completion_mask + 100 + corruption_mask = completion_mask + 50 + batch = { + "inputs": completion_mask + 200, + "completion_mask": completion_mask, + "corruption_mask": corruption_mask, + "targets_loss_mask": targets_loss_mask, + } + + reordered = max_utils.reorder_causal_load_balanced( + batch, + cp_size=2, + reorder_strategy=ReorderStrategy.DUAL_CHUNK_SWAP, + hardware="cpu", + ) + + np.testing.assert_array_equal(reordered["completion_mask"], max_utils.reorder_sequence(completion_mask, 2)) + np.testing.assert_array_equal(reordered["corruption_mask"], max_utils.reorder_sequence(corruption_mask, 2)) + np.testing.assert_array_equal( + reordered["targets_loss_mask"], + max_utils.reorder_sequence(targets_loss_mask, 2), + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py index 28c9f6572a..7f33735c1e 100644 --- a/tests/unit/maxtext_utils_test.py +++ b/tests/unit/maxtext_utils_test.py @@ -1140,13 +1140,22 @@ def test_linen_in_shardings_includes_rng(self): class TestGetShapedBatch(unittest.TestCase): """Tests for get_shaped_batch.""" - def _make_cfg(self, *, enable_diloco=False, use_multimodal=False, use_audio=False): + def _make_cfg( + self, + *, + enable_diloco=False, + use_multimodal=False, + use_audio=False, + training_objective="causal_lm", + ): + """Builds the config subset consumed by get_shaped_batch.""" cfg = MagicMock() cfg.enable_diloco = enable_diloco cfg.global_batch_size_to_load = 4 cfg.max_target_length = 16 cfg.use_multimodal = use_multimodal cfg.use_audio = use_audio + cfg.training_objective = training_objective if enable_diloco: cfg.num_diloco_replicas = 2 return cfg @@ -1169,6 +1178,23 @@ def test_standard_shape(self): expected_shape = (cfg.global_batch_size_to_load, cfg.max_target_length) self.assertEqual(batch["inputs"].shape, expected_shape) + def test_block_diffusion_masks_are_in_shaped_batch(self): + cfg = self._make_cfg(training_objective="block_diffusion") + + batch = maxtext_utils.get_shaped_batch(cfg) + + self.assertEqual(batch["completion_mask"].shape, batch["inputs"].shape) + self.assertEqual(batch["corruption_mask"].shape, batch["inputs"].shape) + self.assertEqual(batch["targets_loss_mask"].shape, batch["inputs"].shape) + self.assertEqual(batch["targets_loss_mask"].dtype, jnp.int32) + + def test_causal_shaped_batch_has_no_diffusion_masks(self): + batch = maxtext_utils.get_shaped_batch(self._make_cfg()) + + self.assertNotIn("completion_mask", batch) + self.assertNotIn("corruption_mask", batch) + self.assertNotIn("targets_loss_mask", batch) + def test_diloco_shape(self): cfg = self._make_cfg(enable_diloco=True) batch = maxtext_utils.get_shaped_batch(cfg) diff --git a/tests/unit/pre_train_loss_mask_test.py b/tests/unit/pre_train_loss_mask_test.py new file mode 100644 index 0000000000..07e4f6200b --- /dev/null +++ b/tests/unit/pre_train_loss_mask_test.py @@ -0,0 +1,264 @@ +# Copyright 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 +# +# http://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. + +"""Tests explicit target-loss masking in the pre-training loss.""" + +# pylint: disable=too-many-positional-arguments + +from dataclasses import dataclass +import unittest +from unittest import mock + +from flax import linen as nn +from flax import nnx +import jax +import jax.numpy as jnp +import numpy as np + +from maxtext.trainers.pre_train import train as pre_train +from maxtext.diffusion import scoring as diffusion_scoring + + +@dataclass +class _Config: + """Configuration subset consumed by pre-training loss_fn.""" + + micro_batch_size_to_train_on: int = 2 + micro_batch_size_to_eval_on: int = 2 + vocab_size: int = 8 + z_loss_multiplier: float = 0.0 + enable_dropout: bool = False + use_multimodal: bool = False + use_indexer: bool = False + indexer_sparse_training: bool = False + indexer_loss_scaling_factor: float = 0.0 + num_vocab_tiling: int = 1 + num_experts: int = 1 + routed_bias: bool = False + routed_bias_update_rate: float = 0.0 + mtp_num_layers: int = 0 + mtp_eval_target_module: int = 0 + use_qk_clip: bool = False + use_tunix_gradient_accumulation: bool = False + gradient_accumulation_steps: int = 1 + shard_mode: int = 0 + debug_sharding: bool = False + weight_sparsity_n: int = 0 + weight_sparsity_m: int = 0 + attention_type: str = "global" + training_objective: str = "causal_lm" + block_diffusion_logit_alignment: str = "same_position" + block_diffusion_canvas_policy: str = "all_masked" + block_diffusion_block_size: int = 4 + + +class _UniformNnxDecoder(nnx.Module): + """Returns uniform logits through the NNX call contract.""" + + def __init__(self, vocab_size): + self.vocab_size = vocab_size + self.mesh = jax.make_mesh((1, 1, 1, 1), ("data", "fsdp", "expert", "context")) + + def __call__( + self, + decoder_input_tokens, + decoder_positions, + decoder_segment_ids=None, + encoder_images=None, + encoder_image_masks=None, + enable_dropout=False, + decoder_target_tokens=None, + decoder_target_mask=None, + ): + del decoder_positions, decoder_segment_ids, encoder_images, encoder_image_masks + del enable_dropout, decoder_target_tokens, decoder_target_mask + return jnp.zeros((*decoder_input_tokens.shape, self.vocab_size), dtype=jnp.float32) + + +class _UniformLinenDecoder(nn.Module): + """Returns uniform logits through the Linen call contract.""" + + vocab_size: int + mesh: object + + @nn.compact + def __call__( + self, + decoder_input_tokens, + decoder_positions, + decoder_segment_ids=None, + encoder_images=None, + encoder_image_masks=None, + enable_dropout=False, + decoder_target_tokens=None, + decoder_target_mask=None, + ): + del decoder_positions, decoder_segment_ids, encoder_images, encoder_image_masks + del enable_dropout, decoder_target_tokens, decoder_target_mask + return jnp.zeros((*decoder_input_tokens.shape, self.vocab_size), dtype=jnp.float32) + + +def _make_data(include_loss_mask=True): + """Builds a batch whose explicit loss mask differs from segmentation.""" + data = { + "inputs": jnp.zeros((2, 4), dtype=jnp.int32), + "inputs_position": jnp.broadcast_to(jnp.arange(4), (2, 4)), + "inputs_segmentation": jnp.ones((2, 4), dtype=jnp.int32), + "targets": jnp.zeros((2, 4), dtype=jnp.int32), + "targets_segmentation": jnp.asarray([[1, 1, 1, 1], [1, 1, 1, 0]], dtype=jnp.int32), + "completion_mask": jnp.asarray([[1, 1, 1, 1], [1, 1, 1, 0]], dtype=jnp.int32), + "corruption_mask": jnp.asarray([[1, 0, 1, 0], [0, 1, 0, 0]], dtype=jnp.int32), + } + if include_loss_mask: + data["targets_loss_mask"] = jnp.asarray([[1, 0, 1, 0], [0, 1, 0, 0]], dtype=jnp.int32) + return data + + +class PreTrainLossMaskTest(unittest.TestCase): + """Checks explicit masks in both Linen and NNX loss branches.""" + + def setUp(self): + super().setUp() + self.config = _Config() + self.per_token_xent = jnp.arange(1, 9, dtype=jnp.float32).reshape(2, 4) + self.per_token_z_loss = self.per_token_xent / 10.0 + + def _cross_entropy_patch(self): + return mock.patch.object( + pre_train.max_utils, + "cross_entropy_with_logits", + return_value=(self.per_token_xent, self.per_token_z_loss), + ) + + def _use_block_diffusion(self): + self.config.attention_type = "block_diffusion" + self.config.training_objective = "block_diffusion" + + def _assert_explicit_mask_result(self, loss, aux): + expected_mask = _make_data()["targets_loss_mask"] != 0 + expected_xent = jnp.sum(self.per_token_xent * expected_mask) + expected_z_loss = jnp.sum(self.per_token_z_loss * expected_mask) / 3.0 + self.assertEqual(int(aux["total_weights"]), 3) + self.assertAlmostEqual(float(aux["xent_sum"]), float(expected_xent)) + self.assertAlmostEqual(float(loss), float(expected_xent / 3.0)) + self.assertAlmostEqual(float(aux["z_loss"]), float(expected_z_loss)) + + def test_nnx_loss_uses_targets_loss_mask(self): + self._use_block_diffusion() + model = _UniformNnxDecoder(self.config.vocab_size) + with self._cross_entropy_patch(): + loss, aux = pre_train.loss_fn(model, self.config, _make_data(), None, None, is_train=True) + + self._assert_explicit_mask_result(loss, aux) + + def test_linen_loss_uses_targets_loss_mask(self): + self._use_block_diffusion() + mesh = jax.make_mesh((1, 1, 1, 1), ("data", "fsdp", "expert", "context")) + model = _UniformLinenDecoder(vocab_size=self.config.vocab_size, mesh=mesh) + data = _make_data() + variables = model.init( + jax.random.key(0), + data["inputs"], + data["inputs_position"], + decoder_segment_ids=data["inputs_segmentation"], + decoder_target_tokens=data["targets"], + decoder_target_mask=data["targets_segmentation"], + ) + with self._cross_entropy_patch(): + loss, aux = pre_train.loss_fn(model, self.config, data, jax.random.key(1), variables, is_train=True) + + self._assert_explicit_mask_result(loss, aux) + + def test_causal_fallback_uses_targets_segmentation(self): + data = _make_data(include_loss_mask=False) + model = _UniformNnxDecoder(self.config.vocab_size) + with self._cross_entropy_patch(): + loss, aux = pre_train.loss_fn(model, self.config, data, None, None, is_train=True) + + expected_mask = data["targets_segmentation"] != 0 + expected_xent = jnp.sum(self.per_token_xent * expected_mask) + self.assertEqual(int(aux["total_weights"]), 7) + self.assertAlmostEqual(float(aux["xent_sum"]), float(expected_xent)) + self.assertAlmostEqual(float(loss), float(expected_xent / 7.0)) + + def test_zero_explicit_mask_has_finite_zero_loss(self): + self._use_block_diffusion() + data = _make_data() + data["targets_loss_mask"] = jnp.zeros_like(data["targets_loss_mask"]) + model = _UniformNnxDecoder(self.config.vocab_size) + with self._cross_entropy_patch(): + loss, aux = pre_train.loss_fn(model, self.config, data, None, None, is_train=True) + + self.assertEqual(int(aux["total_weights"]), 0) + self.assertEqual(float(aux["xent_sum"]), 0.0) + self.assertTrue(bool(jnp.isfinite(loss))) + self.assertEqual(float(loss), 0.0) + + def test_causal_loss_rejects_block_bidirectional_attention(self): + self.config.attention_type = "block_diffusion" + model = _UniformNnxDecoder(self.config.vocab_size) + + with self.assertRaisesRegex(ValueError, "would leak"): + pre_train.loss_fn(model, self.config, _make_data(), None, None, is_train=True) + + def test_block_diffusion_requires_all_explicit_masks(self): + self._use_block_diffusion() + model = _UniformNnxDecoder(self.config.vocab_size) + + for mask_name in ("completion_mask", "corruption_mask", "targets_loss_mask"): + with self.subTest(mask_name=mask_name): + data = _make_data() + del data[mask_name] + with self.assertRaisesRegex(ValueError, mask_name): + pre_train.loss_fn(model, self.config, data, None, None, is_train=True) + + def test_logit_alignment_matches_runtime_contract(self): + logits = jnp.arange(8, dtype=jnp.float32).reshape(1, 4, 2) + + same_position = diffusion_scoring.align_logits_to_targets(logits, "same_position") + shifted = diffusion_scoring.align_logits_to_targets(logits, "shifted") + + np.testing.assert_array_equal(same_position, logits) + np.testing.assert_array_equal(shifted, logits[:, [0, 0, 1, 2], :]) + + def test_shifted_alignment_uses_logical_positions_after_reordering(self): + positions = jnp.asarray([[0, 3, 1, 2]], dtype=jnp.int32) + logits = positions[..., None].astype(jnp.float32) + + shifted = diffusion_scoring.align_logits_to_targets( + logits, + "shifted", + positions=positions, + validity_mask=jnp.ones_like(positions, dtype=jnp.bool_), + ) + + np.testing.assert_array_equal(shifted[..., 0], [[0, 2, 0, 1]]) + + def test_seeded_shifted_canvas_excludes_only_sequence_origin(self): + self.config.attention_type = "block_diffusion" + self.config.training_objective = "block_diffusion" + self.config.block_diffusion_logit_alignment = "shifted" + self.config.block_diffusion_canvas_policy = "seed_and_mask" + data = _make_data() + data["targets_loss_mask"] = jnp.ones_like(data["targets_segmentation"]) + model = _UniformNnxDecoder(self.config.vocab_size) + + with self._cross_entropy_patch(): + _, aux = pre_train.loss_fn(model, self.config, data, None, None, is_train=True) + + self.assertEqual(int(aux["total_weights"]), 5) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/synthetic_data_test.py b/tests/unit/synthetic_data_test.py index 9e29c04a15..f713037a13 100644 --- a/tests/unit/synthetic_data_test.py +++ b/tests/unit/synthetic_data_test.py @@ -15,6 +15,7 @@ """Tests for synthetic data sharding.""" import sys +from types import SimpleNamespace from absl.testing import absltest from absl.testing import parameterized import jax @@ -22,7 +23,7 @@ from jax.experimental import mesh_utils from maxtext.configs import pyconfig -from maxtext.input_pipeline.synthetic_data_processing import SyntheticDataIterator +from maxtext.input_pipeline.synthetic_data_processing import PlaceHolderDataIterator, SyntheticDataIterator from maxtext.utils import sharding as maxtext_sharding from tests.utils.test_helpers import get_test_config_path @@ -75,5 +76,25 @@ def test_synthetic_data_sharding(self, mesh_shape, shard_mode): self.assertEqual(inputs.sharding, expected_sharding) +class PlaceHolderDataTest(absltest.TestCase): + """Tests placeholder batch metadata for non-loading hosts.""" + + def test_block_diffusion_placeholder_has_zero_loss_masks(self): + config = SimpleNamespace( + global_batch_size_to_load=jax.process_count(), + max_target_length=8, + training_objective="block_diffusion", + ) + + batch = next(PlaceHolderDataIterator.get_place_holder_synthetic_data(config)) + + self.assertIn("completion_mask", batch) + self.assertIn("corruption_mask", batch) + self.assertIn("targets_loss_mask", batch) + self.assertFalse(batch["completion_mask"].any()) + self.assertFalse(batch["corruption_mask"].any()) + self.assertFalse(batch["targets_loss_mask"].any()) + + if __name__ == "__main__": absltest.main() diff --git a/tests/unit/tokamax_splash_attention_mask_test.py b/tests/unit/tokamax_splash_attention_mask_test.py index 590e8046b8..b1fe154029 100644 --- a/tests/unit/tokamax_splash_attention_mask_test.py +++ b/tests/unit/tokamax_splash_attention_mask_test.py @@ -20,6 +20,7 @@ from maxtext.kernels.tokamax_splash_attention import splash_attention_mask from maxtext.kernels.tokamax_splash_attention import splash_attention_mask_info +from maxtext.layers.attention_op import BlockCausalMask class TokamaxSplashAttentionMaskTest(absltest.TestCase): @@ -56,6 +57,22 @@ def test_process_causal_mask_keeps_lazy_mask_function(self): self.assertIsNone(mask_info.partial_mask_blocks) np.testing.assert_array_equal(mask_info.q_sequence, np.arange(8, dtype=np.int32)) + def test_process_block_causal_mask_keeps_lazy_mask_function(self): + mask = BlockCausalMask((8, 8), block_size=4) + + mask_info, mask_function = splash_attention_mask_info.process_mask( + mask, + block_shape=(2, 2), + q_seq_shards=1, + kv_seq_shards=1, + return_dynamic_grid=True, + ) + + self.assertIs(mask_function, mask.mask_function) + self.assertIsNotNone(mask_info.q_sequence) + self.assertIsNone(mask_info.partial_mask_blocks) + np.testing.assert_array_equal(mask_info.q_sequence, np.arange(8, dtype=np.int32)) + if __name__ == "__main__": absltest.main()