From cd571e17aabba12db48a6097b32e9ad08d06c641 Mon Sep 17 00:00:00 2001 From: ethannnnnn Date: Tue, 4 Aug 2026 17:11:11 -0700 Subject: [PATCH] [maxtext] Add block-diffusion CFT and SFT Extend the opt-in block-diffusion objective from pre-training to text CFT and completion-only SFT. Preserve clean targets and role-derived completion eligibility through the Hugging Face pipeline, then corrupt only the configured supervision scope. Reject assistant-to-user transitions within one bidirectional diffusion block to prevent future-prompt leakage. Adapt prepared batches to the target-aligned Tunix diffusion contract, retain explicit weighted loss, treat diffusion evaluation as preaveraged, and disable NNX graph caching only for diffusion internal metrics. The causal SFT path remains the default. Draft dependency: temporarily pin the immutable google/tunix#1891 contributor-fork head so CI can exercise the integration. Replace it with the upstream Tunix SHA before marking the PR ready. Test Plan: - 343 passed, 76 platform skips, 3 documented HF integration deselections, 163 subtests - Tunix adapter has 100% statement and branch coverage - Pyink clean and Pylint 10.00/10 - compileall, Yamllint, and git diff --check pass --- .../extra_deps/post_train_github_deps.txt | 2 +- src/maxtext/configs/base.yml | 1 + src/maxtext/configs/types.py | 6 +- .../input_pipeline/hf_data_processing.py | 15 +- .../input_pipeline/input_pipeline_utils.py | 81 +++++++- .../integration/tunix/diffusion_sft.py | 182 ++++++++++++++++++ src/maxtext/trainers/post_train/hooks.py | 6 +- src/maxtext/trainers/post_train/sft/hooks.py | 10 +- .../trainers/post_train/sft/train_sft.py | 31 ++- .../post_training/unit/diffusion_sft_test.py | 180 +++++++++++++++++ tests/post_training/unit/hooks_test.py | 28 +++ tests/post_training/unit/sft_hooks_test.py | 37 ++++ tests/post_training/unit/train_sft_test.py | 98 ++++++++++ tests/unit/configs_value_test.py | 29 ++- tests/unit/hf_data_processing_test.py | 24 ++- tests/unit/input_pipeline_utils_test.py | 160 ++++++++++++++- 16 files changed, 859 insertions(+), 31 deletions(-) create mode 100644 src/maxtext/integration/tunix/diffusion_sft.py create mode 100644 tests/post_training/unit/diffusion_sft_test.py 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."""