diff --git a/src/dependencies/extra_deps/post_train_github_deps.txt b/src/dependencies/extra_deps/post_train_github_deps.txt index cdc56f4053..353584cbfc 100644 --- a/src/dependencies/extra_deps/post_train_github_deps.txt +++ b/src/dependencies/extra_deps/post_train_github_deps.txt @@ -1,3 +1,3 @@ -google-tunix @ https://github.com/google/tunix/archive/c4ec573d29e4c3a3955b348256d464b119c8a6d1.zip +google-tunix @ https://github.com/ethannnnnn/tunix/archive/0ab812a8908e971d3f61d92844f07e12b6ee4806.zip tpu-inference @ https://github.com/vllm-project/tpu-inference/archive/7ecc401e6faefe3f793f3e4a8e6e2dbc49eca868.zip vllm @ git+https://github.com/vllm-project/vllm@0ba2aa35a81dcc3246b26291368b53fa2389c7d7 diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 113160b75f..d7fba03239 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -1032,6 +1032,7 @@ decode_sampling_temperature: 1. eval_start_step: 0 # start eval when train step is >= eval_start_step eval_interval: -1 # the specific number of train step between eval_step eval_steps: -1 # run this number of steps for eval, recommend setting this to prevent error due to running out of evel data +loss_is_preaveraged: false # Whether eval hooks receive a loss already averaged across eval batches target_eval_loss: 0. # early stop once reaching target eval_loss abort_on_nan_loss: true # Check for NaN and abort if found in training loss abort_on_inf_loss: true # Check for Inf and abort if found in training loss diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index a059601f45..a58863e9cf 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1719,6 +1719,10 @@ class TrainingLoop(BaseModel): -1, description="Number of steps to run for each evaluation. -1 runs on entire eval split.", ) + loss_is_preaveraged: bool = Field( + False, + description="Whether the evaluation hook receives a loss already averaged across evaluation batches.", + ) target_eval_loss: float = Field( 0.0, description="If set, training will stop early when this evaluation loss is reached.", @@ -3727,8 +3731,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de 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_sft: - raise ValueError("`training_objective='block_diffusion'` currently supports pre-training only.") if self.use_multimodal or self.use_audio: raise ValueError("`training_objective='block_diffusion'` currently supports text-only training.") valid_model_contracts = { diff --git a/src/maxtext/input_pipeline/hf_data_processing.py b/src/maxtext/input_pipeline/hf_data_processing.py index fd4250c2f3..fb28de61d3 100644 --- a/src/maxtext/input_pipeline/hf_data_processing.py +++ b/src/maxtext/input_pipeline/hf_data_processing.py @@ -48,17 +48,19 @@ def _get_training_objective_transform( shift: bool, use_dpo: bool, use_sft: bool, + completion_only: bool, packing: bool, pad_id: int, bos_token_id: int | None, ) -> input_pipeline_utils.ShiftData | input_pipeline_utils.BlockDiffusionCorruption | None: - """Selects target preparation for causal or block-diffusion pre-training. + """Selects target preparation for causal or block-diffusion training. Args: config: Training configuration containing the objective-specific settings. shift: Whether causal language-model targets should be shifted by one token. use_dpo: Whether the pipeline is preparing direct-preference data. use_sft: Whether the pipeline is preparing supervised fine-tuning data. + completion_only: Whether SFT supervision is restricted to completion tokens. packing: Whether multiple examples are packed into each sequence. pad_id: Token ID used to pad causal language-model examples. bos_token_id: Beginning-of-sequence token ID, or None when unavailable. @@ -68,16 +70,14 @@ def _get_training_objective_transform( Raises: ValueError: If the objective is unsupported or block diffusion is combined with - an incompatible post-training or packing mode. + DPO or packing. """ objective = getattr(config, "training_objective", "causal_lm") if objective == "block_diffusion": - if use_sft: - raise ValueError("This block-diffusion integration currently supports pre-training only.") if use_dpo: - raise ValueError("Block-diffusion pre-training is not compatible with DPO.") + raise ValueError("Block-diffusion training is not compatible with DPO.") if packing: - raise ValueError("Block-diffusion pre-training requires packing=False.") + raise ValueError("Block-diffusion training requires packing=False.") return input_pipeline_utils.BlockDiffusionCorruption( block_size=config.causal_block_size, mask_id=config.block_diffusion_mask_id, @@ -85,6 +85,7 @@ def _get_training_objective_transform( logit_alignment=config.block_diffusion_logit_alignment, canvas_policy=config.block_diffusion_canvas_policy, axis=1, + completion_only=bool(use_sft and completion_only), ) if objective != "causal_lm": raise ValueError(f"Unsupported training objective: {objective}") @@ -389,6 +390,7 @@ def preprocessing_pipeline( completion_only=sft_train_on_completion_only, max_target_length=max_target_length, unk_id=pad_id, + training_objective=getattr(config, "training_objective", "causal_lm"), ) ) data_column_names = ("inputs", "targets") @@ -424,6 +426,7 @@ def preprocessing_pipeline( 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, diff --git a/src/maxtext/input_pipeline/input_pipeline_utils.py b/src/maxtext/input_pipeline/input_pipeline_utils.py index bf993b25be..dfe12e0f31 100644 --- a/src/maxtext/input_pipeline/input_pipeline_utils.py +++ b/src/maxtext/input_pipeline/input_pipeline_utils.py @@ -329,33 +329,51 @@ 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``. Block diffusion keeps + clean same-position targets and emits role-derived completion metadata. """ - 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, + training_objective="causal_lm", + ): + if training_objective not in ("causal_lm", "block_diffusion"): + raise ValueError(f"Unsupported training objective: {training_objective}") 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.training_objective = training_objective def map(self, element): """ Maps a single dataset element to an SFT training instance. 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`. + - If `self.completion_only` is `True`, causal SFT masks the prompt portion + with `self.unk_id`; block-diffusion SFT keeps clean targets and records + completion eligibility separately. - If `self.completion_only` is `False`, the target sequence is identical to the input sequence. """ - inputs, targets = [], [] + is_block_diffusion = self.training_objective == "block_diffusion" + 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 is_block_diffusion 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 is_block_diffusion: + output["completion_mask"] = np.asarray(completion_mask[: self.max_target_length], dtype=np.int32) + return output @dataclasses.dataclass @@ -854,7 +872,7 @@ def map( self.config is not None and getattr(self.config, "training_objective", "causal_lm") == "block_diffusion" ) 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.") @@ -887,7 +905,7 @@ def map( element["images"] = self._pad_image_and_mask(element["images"]) # pyrefly: ignore[bad-argument-type] - elif preserve_pad_valued_tokens and key.endswith(("_segmentation", "_position")): + elif key == "completion_mask" or (preserve_pad_valued_tokens and key.endswith(("_segmentation", "_position"))): 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] @@ -1026,6 +1044,7 @@ class BlockDiffusionCorruption(grain.RandomMapTransform): logit_alignment: str = "same_position" canvas_policy: str = "all_masked" axis: int = 1 + completion_only: bool = False def random_map(self, element, rng: np.random.Generator): """Corrupts inputs while preserving clean targets and input metadata.""" @@ -1037,9 +1056,51 @@ def random_map(self, element, rng: np.random.Generator): "inputs, targets, and targets_segmentation must have identical shapes, got " f"{inputs.shape}, {targets.shape}, and {targets_segmentation.shape}" ) + validity_mask = targets_segmentation != 0 + raw_completion_mask = element.get("completion_mask") + if raw_completion_mask is None: + if self.completion_only: + if "completion_mask" not in element: + raise ValueError("completion-only block diffusion requires an explicit completion_mask") + raise ValueError("completion-only block diffusion requires a non-None completion_mask") + completion_mask = validity_mask + else: + completion_mask = np.asarray(raw_completion_mask) != 0 + if completion_mask.shape != inputs.shape: + raise ValueError(f"completion_mask must match inputs shape; received {completion_mask.shape} and {inputs.shape}") + if np.any(completion_mask & ~validity_mask): + raise ValueError("completion_mask must be a subset of valid target positions") + can_validate_block_order = ( + self.completion_only + and self.block_size > 0 + and inputs.ndim > 0 + and -inputs.ndim <= self.axis < inputs.ndim + and inputs.shape[self.axis % inputs.ndim] > 0 + ) + if can_validate_block_order: + axis = self.axis % inputs.ndim + sequence_length = inputs.shape[axis] + block_count = (sequence_length + self.block_size - 1) // self.block_size + padded_length = block_count * self.block_size + + def _as_blocks(mask): + rows = np.moveaxis(mask, axis, -1).reshape(-1, sequence_length) + padded = np.pad(rows, ((0, 0), (0, padded_length - sequence_length))) + return padded.reshape(rows.shape[0], block_count, self.block_size) + + validity_blocks = _as_blocks(validity_mask) + completion_blocks = _as_blocks(completion_mask) + seen_completion = np.maximum.accumulate(completion_blocks, axis=-1) + if np.any(seen_completion & validity_blocks & ~completion_blocks): + raise ValueError( + "completion-only block diffusion cannot place a prompt token after " + "a completion token in the same diffusion block; split the example " + "or align the next prompt to a block boundary" + ) + supervision_mask = completion_mask if self.completion_only else validity_mask result = block_diffusion_corruption.corrupt_tokens( inputs, - targets_segmentation != 0, + supervision_mask, rng, block_size=self.block_size, mask_id=self.mask_id, diff --git a/src/maxtext/integration/tunix/diffusion_sft.py b/src/maxtext/integration/tunix/diffusion_sft.py new file mode 100644 index 0000000000..b487eea076 --- /dev/null +++ b/src/maxtext/integration/tunix/diffusion_sft.py @@ -0,0 +1,182 @@ +# 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.block_diffusion import target_alignment +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): + """Returns a NumPy view only for values already resident on the host.""" + if isinstance(value, (jax.Array, jax.core.Tracer, jax.ShapeDtypeStruct)): + return None + return np.asarray(value) + + +def _validate_batch_masks( + positions, + targets_segmentation, + completion_mask, + corruption_mask, + loss_weights, + *, + alignment, + block_size, + completion_only, +): + """Validates the prepared block-diffusion supervision contract.""" + shapes = { + "positions": tuple(positions.shape), + "targets_segmentation": tuple(targets_segmentation.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, + targets_segmentation, + completion_mask, + corruption_mask, + loss_weights, + ) + ] + if any(value is None for value in concrete): + return + concrete_positions = np.asarray(concrete[0], dtype=np.int32) + validity = np.asarray(concrete[1]) != 0 + completion, corruption = (np.asarray(value, dtype=bool) for value in concrete[2:4]) + weights = np.asarray(concrete[4], dtype=np.float32) + 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 completion_only: + block_ids = concrete_positions // block_size + unseen_block = np.iinfo(block_ids.dtype).min + seen_completion_blocks = np.maximum.accumulate(np.where(validity & completion, block_ids, unseen_block), axis=-1) + if np.any(validity & ~completion & (seen_completion_blocks == block_ids)): + raise ValueError( + "completion-only block diffusion cannot place a prompt token after " + "a completion token in the same diffusion block; split the example " + "or align the next prompt to a block boundary" + ) + supervision = completion if completion_only else validity + if np.any(corruption & ~supervision): + raise ValueError("corruption_mask must be a subset of the configured supervision scope") + if np.any(weighted & ~supervision): + raise ValueError("targets_loss_mask must be a subset of the configured supervision scope") + 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.causal_block_size) + completion_only = bool(config.sft_train_on_completion_only) + + def adapt(raw_batch: Mapping[str, Any]) -> diffusion_types.DiffusionTokenBatch: + unavailable = sorted(name for name in _REQUIRED_FIELDS if name not in raw_batch or raw_batch[name] is None) + if unavailable: + raise ValueError("block-diffusion SFT requires non-None batch fields; " f"unavailable {unavailable}") + _validate_batch_masks( + raw_batch["targets_position"], + raw_batch["targets_segmentation"], + raw_batch["completion_mask"], + raw_batch["corruption_mask"], + raw_batch["targets_loss_mask"], + alignment=alignment, + block_size=block_size, + completion_only=completion_only, + ) + 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) + supervision_mask = completion_mask if completion_only else validity_mask + allowed = corruption_mask + if alignment == "shifted": + allowed |= (positions > 0) & (positions % block_size == 0) + loss_weights = jnp.where(validity_mask & supervision_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 target_alignment.align_logits_to_targets( + logits, + alignment, + model_inputs["target_positions"], + model_inputs["target_segmentation"] != 0, + ) + + return logits_fn diff --git a/src/maxtext/trainers/post_train/hooks.py b/src/maxtext/trainers/post_train/hooks.py index bc8666a1c5..3152131943 100644 --- a/src/maxtext/trainers/post_train/hooks.py +++ b/src/maxtext/trainers/post_train/hooks.py @@ -162,7 +162,7 @@ def on_eval_step_end(self, train_ctx: peft_trainer.PeftTrainer, eval_loss: float self.eval_metadata["eval_step_count"] != 0 ), "BaseTrainingHooks.on_eval_step_start() must be called before BaseTrainingHooks.on_eval_step_end()" - avg_loss = eval_loss / self.eval_metadata["eval_step_count"] + avg_loss = eval_loss if self.eval_loss_is_preaveraged() else eval_loss / self.eval_metadata["eval_step_count"] metrics = { "scalar": { "eval/total_loss": eval_loss, @@ -195,6 +195,10 @@ def on_eval_step_end(self, train_ctx: peft_trainer.PeftTrainer, eval_loss: float if avg_loss <= self.config.target_eval_loss: raise exceptions.StopTraining(f"Target loss {self.config.target_eval_loss=} is achieved.") + def eval_loss_is_preaveraged(self) -> bool: + """Whether Tunix supplies a cross-batch weighted mean to the eval hook.""" + return False + @abc.abstractmethod def get_total_weights(self, batch) -> jax.Array: """Calculate the number of non-padded tokens in the batch.""" diff --git a/src/maxtext/trainers/post_train/sft/hooks.py b/src/maxtext/trainers/post_train/sft/hooks.py index 9f30a3d6d8..f2648fe164 100644 --- a/src/maxtext/trainers/post_train/sft/hooks.py +++ b/src/maxtext/trainers/post_train/sft/hooks.py @@ -29,7 +29,15 @@ 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) + loss_weights = batch.get("targets_loss_mask") + if loss_weights is None: + loss_weights = batch["targets_segmentation"] + return jnp.sum(loss_weights != 0) + + @override + def eval_loss_is_preaveraged(self) -> bool: + """Whether the configured objective supplies a preaveraged eval loss.""" + return bool(self.config.loss_is_preaveraged or self.config.training_objective == "block_diffusion") 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 fb2926e67f..ab98650352 100644 --- a/src/maxtext/trainers/post_train/sft/train_sft.py +++ b/src/maxtext/trainers/post_train/sft/train_sft.py @@ -50,10 +50,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, @@ -277,6 +279,24 @@ def loss_func( return trainer +def configure_training_objective(trainer, mt_config): + """Configures causal or target-aligned diffusion SFT.""" + 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: @@ -289,6 +309,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): @@ -315,10 +336,10 @@ def setup_trainer_state(mt_config, goodput_recorder=None): nnx.pop(model, nnx.Intermediate) if mt_config.lora.lora_restore_path: lora_utils.restore_lora_from_path(model, mt_config) - trainer = MaxTextPeftTrainer(model, optimizer, tunix_config) + trainer = _create_trainer(model, optimizer, tunix_config, 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 @@ -326,9 +347,9 @@ def setup_trainer_state(mt_config, goodput_recorder=None): def train_model(mt_config, trainer, mesh): """Runs the SFT training loop in Tunix.""" with jax.set_mesh(mesh), nn_partitioning.axis_rules(mt_config.logical_axis_rules): - # Disable NNX graph caching for MoE models (where experts > 1) to allow - # necessary dynamic metadata synchronization during forward passes (e.g., in jax.lax.scan). - enable_nnx_cache = mt_config.num_experts <= 1 + is_block_diffusion = getattr(mt_config, "training_objective", "causal_lm") == "block_diffusion" + records_internal_metrics = bool(getattr(mt_config, "record_internal_nn_metrics", False)) + enable_nnx_cache = mt_config.num_experts <= 1 and not (is_block_diffusion and records_internal_metrics) trainer.train( trainer.data_hooks.train_data_iterator, 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..31125f0fd5 --- /dev/null +++ b/tests/post_training/unit/diffusion_sft_test.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. + +"""Tests for the MaxText-to-Tunix diffusion SFT adapter.""" + +from types import SimpleNamespace +from unittest import mock + +import jax.numpy as jnp +import numpy as np +import pytest + +from maxtext.diffusion.block_diffusion import target_alignment +from maxtext.integration.tunix import diffusion_sft + + +pytestmark = [pytest.mark.post_training] + + +def _config(alignment="shifted", completion_only=True): + return SimpleNamespace( + causal_block_size=4, + block_diffusion_logit_alignment=alignment, + sft_train_on_completion_only=completion_only, + enable_dropout=False, + ) + + +def _raw_batch(array_module=np): + positions = array_module.arange(8, dtype=np.int32)[None, :] + segmentation = array_module.ones((1, 8), dtype=np.int32) + return { + "inputs": array_module.asarray([[10, 11, 99, 99, 14, 99, 99, 99]], dtype=np.int32), + "inputs_position": positions, + "inputs_segmentation": segmentation, + "targets": array_module.asarray([[10, 11, 12, 13, 14, 15, 16, 17]], dtype=np.int32), + "targets_position": positions, + "targets_segmentation": segmentation, + "completion_mask": array_module.asarray([[0, 0, 1, 1, 1, 1, 1, 1]], dtype=np.int32), + "corruption_mask": array_module.asarray([[0, 0, 1, 1, 0, 1, 1, 1]], dtype=np.int32), + "targets_loss_mask": array_module.asarray([[0, 0, 1, 1, 1, 1, 1, 1]], dtype=np.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"][0, 3] = 0 + + with pytest.raises(ValueError, match="corrupted targets or shifted block anchors"): + diffusion_sft.create_batch_adapter(_config())(raw) + + +@pytest.mark.parametrize("field", ["targets_position", "targets_loss_mask"]) +def test_batch_adapter_rejects_none_required_field(field): + raw = _raw_batch() + raw[field] = None + + with pytest.raises(ValueError, match=field): + diffusion_sft.create_batch_adapter(_config())(raw) + + +def test_batch_adapter_preserves_full_token_cft_supervision(): + raw = _raw_batch() + raw["corruption_mask"] = np.ones_like(raw["corruption_mask"]) + raw["targets_loss_mask"] = np.ones_like(raw["targets_loss_mask"]) + + batch = diffusion_sft.create_batch_adapter(_config(alignment="same_position", completion_only=False))(raw) + + np.testing.assert_array_equal(batch.loss_weights, jnp.ones((1, 8))) + + +def test_completion_only_adapter_rejects_prompt_supervision(): + raw = _raw_batch() + raw["corruption_mask"][0, 0] = 1 + raw["targets_loss_mask"][0, 0] = 1 + + with pytest.raises(ValueError, match="configured supervision scope"): + diffusion_sft.create_batch_adapter(_config())(raw) + + +def test_completion_only_adapter_rejects_future_prompt_in_block(): + raw = _raw_batch() + raw["completion_mask"][0, 5] = 0 + raw["corruption_mask"][0, 5] = 0 + raw["targets_loss_mask"][0, 5] = 0 + + with pytest.raises(ValueError, match="prompt token after a completion token"): + diffusion_sft.create_batch_adapter(_config())(raw) + + +def test_concrete_numpy_does_not_materialize_jax_arrays(): + value = jnp.ones((1, 2)) + + with mock.patch.object(diffusion_sft.np, "asarray") as asarray: + assert diffusion_sft._concrete_numpy(value) is None # pylint: disable=protected-access + + asarray.assert_not_called() + + +def test_batch_adapter_accepts_jax_arrays_without_eager_value_validation(): + batch = diffusion_sft.create_batch_adapter(_config())(_raw_batch(jnp)) + + np.testing.assert_array_equal(batch.target_ids, _raw_batch()["targets"]) + + +def test_batch_adapter_rejects_mismatched_mask_shapes(): + raw = _raw_batch() + raw["completion_mask"] = raw["completion_mask"][:, :-1] + + with pytest.raises(ValueError, match="must have identical shapes"): + diffusion_sft.create_batch_adapter(_config())(raw) + + +@pytest.mark.parametrize("invalid_weight", [-1.0, np.nan]) +def test_batch_adapter_rejects_invalid_loss_weights(invalid_weight): + raw = _raw_batch() + raw["targets_loss_mask"] = raw["targets_loss_mask"].astype(np.float32) + raw["targets_loss_mask"][0, 2] = invalid_weight + + with pytest.raises(ValueError, match="finite nonnegative weights"): + diffusion_sft.create_batch_adapter(_config())(raw) + + +def test_batch_adapter_rejects_completion_outside_validity(): + raw = _raw_batch() + raw["targets_segmentation"][0, 2] = 0 + + with pytest.raises(ValueError, match="subset of valid target positions"): + diffusion_sft.create_batch_adapter(_config())(raw) + + +def test_batch_adapter_rejects_loss_weight_outside_supervision(): + raw = _raw_batch() + raw["targets_loss_mask"][0, 0] = 1 + + with pytest.raises(ValueError, match="targets_loss_mask must be a subset"): + 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 = target_alignment.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/hooks_test.py b/tests/post_training/unit/hooks_test.py index f34924eb7a..4321a1c3ba 100644 --- a/tests/post_training/unit/hooks_test.py +++ b/tests/post_training/unit/hooks_test.py @@ -45,6 +45,13 @@ def get_total_weights(self, batch) -> jax.Array: return np.sum(batch["targets_segmentation"] != 0) +class PreaveragedDummyTrainingHooks(DummyTrainingHooks): + """Models an objective that passes a weighted eval mean.""" + + def eval_loss_is_preaveraged(self) -> bool: + return True + + class BaseHooksTest(unittest.TestCase): def setUp(self): @@ -118,6 +125,27 @@ def test_training_hooks_for_eval_step(self): self.assertAlmostEqual(metrics["eval/avg_perplexity"], np.exp(5.0), places=2) self.assertEqual(metrics["eval/total_weights"], jax.device_count() * self.config.max_target_length * total_eval_steps) + def test_preaveraged_eval_loss_is_not_divided_again(self): + training_hooks = PreaveragedDummyTrainingHooks( + self.config, + self.mesh, + self.learning_rate_schedule, + goodput_recorder=None, + ) + training_hooks.metric_logger = MetricLogger(self.config, self.learning_rate_schedule) + training_hooks.metric_logger.metadata = defaultdict(float) + self.mock_train_ctx.data_hooks.eval_batch = self.expected_batch + self.mock_train_ctx.train_steps = 0 + for _ in range(2): + training_hooks.on_eval_step_start(self.mock_train_ctx) + + training_hooks.on_eval_step_end(self.mock_train_ctx, eval_loss=3.0) + + metrics = self._read_logged_metrics(num_expected=1)[0] + self.assertAlmostEqual(metrics["eval/total_loss"], 3.0) + self.assertAlmostEqual(metrics["eval/avg_loss"], 3.0) + self.assertAlmostEqual(metrics["eval/avg_perplexity"], np.exp(3.0), places=2) + def test_on_train_end_asserts_if_on_train_start_not_called(self): with self.assertRaises(AssertionError): self.training_hooks.on_train_end(self.mock_train_ctx) diff --git a/tests/post_training/unit/sft_hooks_test.py b/tests/post_training/unit/sft_hooks_test.py index e8fa00b123..155d24b0fc 100644 --- a/tests/post_training/unit/sft_hooks_test.py +++ b/tests/post_training/unit/sft_hooks_test.py @@ -14,6 +14,7 @@ """Tests for training and data loading hooks for SFT""" +from types import SimpleNamespace import unittest from unittest.mock import MagicMock, patch @@ -74,6 +75,42 @@ 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) + + def test_sft_training_hooks_falls_back_from_none_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.array([[1, 1, 0], [1, 0, 0]]), + "targets_loss_mask": None, + } + + total_weights = training_hooks.get_total_weights(batch) + + self.assertEqual(total_weights, 3) + + def test_preaveraged_eval_loss_is_config_driven(self): + training_hooks = object.__new__(sft_hooks.SFTTrainingHooks) + + training_hooks.config = SimpleNamespace(loss_is_preaveraged=False, training_objective="causal_lm") + self.assertFalse(training_hooks.eval_loss_is_preaveraged()) + + training_hooks.config = SimpleNamespace(loss_is_preaveraged=True, training_objective="causal_lm") + self.assertTrue(training_hooks.eval_loss_is_preaveraged()) + + training_hooks.config = SimpleNamespace(loss_is_preaveraged=False, training_objective="block_diffusion") + self.assertTrue(training_hooks.eval_loss_is_preaveraged()) + if __name__ == "__main__": unittest.main() diff --git a/tests/post_training/unit/train_sft_test.py b/tests/post_training/unit/train_sft_test.py index 3e71ba6273..79bd4b39b2 100644 --- a/tests/post_training/unit/train_sft_test.py +++ b/tests/post_training/unit/train_sft_test.py @@ -44,6 +44,83 @@ def test_validate_config_invalid_offload(self): with self.assertRaisesRegex(ValueError, "optimizer_memory_host_offload=True is not supported"): train_sft.validate_config(config) + 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) + + def test_validate_config_accepts_preaveraged_diffusion_evaluation(self): + config = SimpleNamespace( + optimizer_memory_host_offload=False, + training_objective="block_diffusion", + gradient_accumulation_steps=1, + eval_interval=10, + loss_is_preaveraged=True, + ) + + train_sft.validate_config(config) + + 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() + + 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") + + def test_causal_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) # pylint: disable=protected-access + 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) + + def test_diffusion_objective_uses_weighted_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) # pylint: disable=protected-access + + self.assertIs(trainer, mock.sentinel.trainer) + trainer_type.assert_called_once_with("model", "optimizer", "config") + def test_train_model_caching_moe(self): """Test that NNX graph caching is disabled for MoE models (num_experts > 1).""" mt_config = SimpleNamespace( @@ -103,6 +180,27 @@ def test_maxtext_peft_trainer_train_step_signature(self): params = list(sig.parameters.keys()) self.assertEqual(params, ["model", "optimizer", "grad_accumulator", "inputs", "is_update_step"]) + def test_train_model_disables_cache_for_diffusion_internal_metrics(self): + mt_config = SimpleNamespace( + logical_axis_rules=[], + num_experts=1, + training_objective="block_diffusion", + record_internal_nn_metrics=True, + ) + trainer = mock.MagicMock() + trainer.data_hooks.train_data_iterator = "train_iter" + trainer.data_hooks.eval_data_iterator = "eval_iter" + mesh = mock.MagicMock() + + with mock.patch("jax.set_mesh"): + train_sft.train_model(mt_config, trainer, mesh) + + trainer.train.assert_called_once_with( + "train_iter", + "eval_iter", + cache_nnx_graph=False, + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index 1b4a15a913..42999706fa 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -631,6 +631,7 @@ def test_default_attention_remains_global(self): self.assertEqual(config.attention_type, "global") self.assertEqual(config.training_objective, "causal_lm") self.assertEqual(config.block_diffusion_mask_id, -1) + self.assertFalse(config.loss_is_preaveraged) def test_block_diffusion_pretraining_config(self): config = pyconfig.initialize( @@ -687,7 +688,6 @@ def test_block_diffusion_pretraining_rejects_incompatible_config(self): ({"num_vocab_tiling": 2}, "vocabulary tiling"), ({"dataset_type": "grain"}, "dataset_type='hf'"), ({"use_dpo": True}, "DPO"), - ({"use_sft": True}, "pre-training only"), ({"use_multimodal": True}, "text-only"), ({"block_diffusion_logit_alignment": "shifted"}, "seed_and_mask"), ({"block_diffusion_canvas_policy": "seed_and_mask"}, "same_position/all_masked"), @@ -707,6 +707,33 @@ def test_block_diffusion_pretraining_rejects_incompatible_config(self): with self.assertRaisesRegex((ValueError, pydantic.ValidationError), expected_regex): pyconfig.initialize(argv) + def test_block_diffusion_sft_config(self): + config = pyconfig.initialize( + [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "steps=1", + "training_objective=block_diffusion", + "attention=dot_product", + "attention_type=block_diffusion", + "causal_block_size=8", + "block_diffusion_mask_id=100", + "vocab_size=256", + "packing=False", + "dataset_type=hf", + "hf_path=parquet", + "use_sft=True", + "sft_train_on_completion_only=True", + "loss_is_preaveraged=True", + "hardware=cpu", + ] + ) + + self.assertTrue(config.use_sft) + self.assertTrue(config.sft_train_on_completion_only) + self.assertTrue(config.loss_is_preaveraged) + def test_shifted_block_diffusion_requires_seeded_canvas(self): config = pyconfig.initialize( [ diff --git a/tests/unit/hf_data_processing_test.py b/tests/unit/hf_data_processing_test.py index 3b2c7bf738..9162df8a86 100644 --- a/tests/unit/hf_data_processing_test.py +++ b/tests/unit/hf_data_processing_test.py @@ -190,6 +190,7 @@ def test_default_objective_keeps_next_token_shift(self): shift=True, use_dpo=False, use_sft=False, + completion_only=False, packing=False, pad_id=0, bos_token_id=1, @@ -204,6 +205,7 @@ def test_causal_objective_without_shift_has_no_transform(self): shift=False, use_dpo=False, use_sft=False, + completion_only=False, packing=False, pad_id=0, bos_token_id=1, @@ -218,6 +220,7 @@ def test_unknown_objective_is_rejected(self): shift=False, use_dpo=False, use_sft=False, + completion_only=False, packing=False, pad_id=0, bos_token_id=1, @@ -229,6 +232,7 @@ def test_block_diffusion_replaces_next_token_shift(self): shift=True, use_dpo=False, use_sft=False, + completion_only=False, packing=False, pad_id=0, bos_token_id=1, @@ -240,6 +244,7 @@ def test_block_diffusion_replaces_next_token_shift(self): self.assertEqual(transform.min_noise, 0.05) self.assertEqual(transform.logit_alignment, "shifted") self.assertEqual(transform.canvas_policy, "seed_and_mask") + self.assertFalse(transform.completion_only) def test_preprocessing_pipeline_installs_block_diffusion_transform(self): operations = self._pipeline_operations(self._block_diffusion_config(), shift=True) @@ -261,18 +266,33 @@ def test_preprocessing_pipeline_omits_disabled_causal_shift(self): ) ) - def test_block_diffusion_rejects_packing_and_post_training_modes(self): + def test_block_diffusion_sft_uses_completion_scope(self): + transform = hf_data_processing._get_training_objective_transform( # pylint: disable=protected-access + self._block_diffusion_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.assertTrue(transform.completion_only) + + def test_block_diffusion_rejects_packing_and_dpo(self): base_args = { "shift": True, "use_dpo": False, "use_sft": False, + "completion_only": False, "packing": False, "pad_id": 0, "bos_token_id": 1, } cases = ( ({"packing": True}, "packing=False"), - ({"use_sft": True}, "pre-training only"), ({"use_dpo": True}, "not compatible with DPO"), ) for overrides, expected_message in cases: diff --git a/tests/unit/input_pipeline_utils_test.py b/tests/unit/input_pipeline_utils_test.py index 46ae9eca47..8c627d2966 100644 --- a/tests/unit/input_pipeline_utils_test.py +++ b/tests/unit/input_pipeline_utils_test.py @@ -20,7 +20,12 @@ import numpy as np -from maxtext.input_pipeline.input_pipeline_utils import BlockDiffusionCorruption, compute_file_sharding, PadOrTrimToMaxLength +from maxtext.input_pipeline.input_pipeline_utils import ( + BlockDiffusionCorruption, + compute_file_sharding, + PadOrTrimToMaxLength, + SFTPromptMasking, +) class BlockDiffusionPaddingTest(unittest.TestCase): @@ -39,7 +44,15 @@ def test_corruption_dataclass_serializes_configuration(self): self.assertEqual( tuple(field.name for field in dataclasses.fields(transform)), - ("block_size", "mask_id", "min_noise", "logit_alignment", "canvas_policy", "axis"), + ( + "block_size", + "mask_id", + "min_noise", + "logit_alignment", + "canvas_policy", + "axis", + "completion_only", + ), ) self.assertNotEqual(dataclasses.asdict(transform), dataclasses.asdict(other)) self.assertIn("block_size=4", repr(transform)) @@ -114,6 +127,149 @@ def test_corruption_rejects_mismatched_batch_shapes(self): ) +class BlockDiffusionSFTInputTest(unittest.TestCase): + """Checks completion roles remain separate from corruption and validity.""" + + def _prepared_batch(self): + """Builds a block-aligned conversation with two completion spans.""" + clean = SFTPromptMasking( + text_column_name="messages", + completion_only=True, + max_target_length=8, + unk_id=0, + training_objective="block_diffusion", + ).map( + { + "messages": [[11, 12], [21, 22], [31], [41]], + "is_prompt": [True, False, True, False], + } + ) + return PadOrTrimToMaxLength( + max_length=8, + pad_id=7, + config=SimpleNamespace(training_objective="block_diffusion"), + ).map(clean) + + def test_role_mask_preserves_clean_targets_and_padding(self): + batch = self._prepared_batch() + + np.testing.assert_array_equal(batch["inputs"], batch["targets"]) + np.testing.assert_array_equal(batch["completion_mask"], [0, 0, 1, 1, 0, 1, 0, 0]) + np.testing.assert_array_equal(batch["targets_segmentation"], [1, 1, 1, 1, 1, 1, 0, 0]) + self.assertNotIn("completion_mask_segmentation", batch) + self.assertNotIn("completion_mask_position", batch) + + def test_prompt_masking_rejects_unsupported_objective(self): + with self.assertRaisesRegex(ValueError, "Unsupported training objective"): + SFTPromptMasking( + text_column_name="messages", + completion_only=True, + max_target_length=8, + training_objective="unsupported", + ) + + def test_completion_only_corruption_never_supervises_prompt(self): + batch = self._prepared_batch() + output = BlockDiffusionCorruption( + block_size=4, + mask_id=99, + min_noise=1.0, + completion_only=True, + axis=0, + ).random_map(batch, np.random.default_rng(0)) + + completion = batch["completion_mask"] != 0 + self.assertFalse(output["corruption_mask"][~completion].any()) + self.assertFalse(output["targets_loss_mask"][~completion].any()) + np.testing.assert_array_equal(output["targets"], batch["targets"]) + + def test_completion_only_corruption_requires_role_mask(self): + batch = self._prepared_batch() + del batch["completion_mask"] + + with self.assertRaisesRegex(ValueError, "explicit completion_mask"): + BlockDiffusionCorruption( + block_size=4, + mask_id=99, + completion_only=True, + axis=0, + ).random_map(batch, np.random.default_rng(0)) + + def test_completion_only_corruption_rejects_none_role_mask(self): + batch = self._prepared_batch() + batch["completion_mask"] = None + + with self.assertRaisesRegex(ValueError, "non-None completion_mask"): + BlockDiffusionCorruption( + block_size=4, + mask_id=99, + completion_only=True, + axis=0, + ).random_map(batch, np.random.default_rng(0)) + + def test_completion_only_corruption_rejects_mismatched_role_mask_shape(self): + batch = self._prepared_batch() + batch["completion_mask"] = batch["completion_mask"][:-1] + + with self.assertRaisesRegex(ValueError, "completion_mask must match inputs shape"): + BlockDiffusionCorruption( + block_size=4, + mask_id=99, + completion_only=True, + axis=0, + ).random_map(batch, np.random.default_rng(0)) + + def test_completion_only_corruption_rejects_role_mask_on_padding(self): + batch = self._prepared_batch() + batch["completion_mask"][6] = 1 + + with self.assertRaisesRegex(ValueError, "subset of valid target positions"): + BlockDiffusionCorruption( + block_size=4, + mask_id=99, + completion_only=True, + axis=0, + ).random_map(batch, np.random.default_rng(0)) + + def test_completion_only_corruption_rejects_future_prompt_in_block(self): + batch = self._prepared_batch() + batch["completion_mask"] = np.asarray([0, 0, 1, 1, 1, 0, 0, 0], dtype=np.int32) + + with self.assertRaisesRegex(ValueError, "prompt token after a completion token"): + BlockDiffusionCorruption( + block_size=4, + mask_id=99, + completion_only=True, + axis=0, + ).random_map(batch, np.random.default_rng(0)) + + def test_full_token_corruption_treats_none_role_mask_as_absent(self): + batch = self._prepared_batch() + batch["completion_mask"] = None + + output = BlockDiffusionCorruption( + block_size=4, + mask_id=99, + min_noise=1.0, + axis=0, + ).random_map(batch, np.random.default_rng(0)) + + self.assertTrue(output["corruption_mask"][:6].all()) + self.assertFalse(output["corruption_mask"][6:].any()) + + def test_causal_sft_contract_is_unchanged(self): + batch = SFTPromptMasking( + text_column_name="messages", + completion_only=True, + max_target_length=4, + unk_id=0, + ).map({"messages": [[11, 12], [21, 22]], "is_prompt": [True, False]}) + + np.testing.assert_array_equal(batch["inputs"], [11, 12, 21, 22]) + np.testing.assert_array_equal(batch["targets"], [0, 0, 21, 22]) + self.assertNotIn("completion_mask", batch) + + class ComputeFileShardingNormalCaseTest(unittest.TestCase): """file_count >= host_count: disjoint file subsets, no row sharding."""