From aadaaa8abcff2d6cdececd51b31bfcbf89e68b06 Mon Sep 17 00:00:00 2001 From: ethannnnnn Date: Tue, 4 Aug 2026 16:59:52 -0700 Subject: [PATCH] [maxtext] Add block-diffusion pre-training Wire the block-diffusion primitives into the opt-in text pre-training path. HF batches keep clean same-position targets, corrupt valid tokens per block, and carry separate corruption and loss masks. The loss aligns logits according to the configured model contract and normalizes only supervised positions. Causal LM remains the default; packing, SFT, DPO, MTP, vocabulary tiling, and multimodal input fail validation for this initial scope. Test Plan: - corruption adapter and nonzero-pad metadata tests - Linen and NNX loss-mask tests - objective/config compatibility tests - positive-weight causal and zero-weight native gradient-accumulation tests --- src/maxtext/configs/base.yml | 6 +- src/maxtext/configs/types.py | 57 ++++ .../input_pipeline/hf_data_processing.py | 66 ++++- .../input_pipeline/input_pipeline_utils.py | 66 ++++- src/maxtext/trainers/pre_train/train.py | 57 +++- src/maxtext/utils/gradient_accumulation.py | 22 +- src/maxtext/utils/max_utils.py | 2 + src/maxtext/utils/maxtext_utils.py | 3 + tests/unit/configs_value_test.py | 102 +++++++ tests/unit/gradient_accumulation_nnx_test.py | 95 ++++++- tests/unit/hf_data_processing_test.py | 159 +++++++++++ tests/unit/input_pipeline_utils_test.py | 80 +++++- tests/unit/max_utils_test.py | 24 ++ tests/unit/maxtext_utils_test.py | 28 +- tests/unit/pre_train_loss_mask_test.py | 251 ++++++++++++++++++ 15 files changed, 994 insertions(+), 24 deletions(-) create mode 100644 tests/unit/pre_train_loss_mask_test.py diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 7a725fc4ca..ba50f4c7b3 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -806,6 +806,11 @@ olmo_apply_ngram_filter: true # mask instances with repetitive n-grams (OLMo-cor # Training loop steps: 150_001 # If set to -1 then will inherit value from learning_rate_schedule_steps log_period: 100 # The frequency of Tensorboard flush, gcs metrics writing, and managed profiler metrics updating. +training_objective: 'causal_lm' # Supported objectives: causal_lm, block_diffusion +block_diffusion_mask_id: -1 # Tokenizer mask-token id; required for training_objective='block_diffusion' +block_diffusion_min_noise: 0.001 # Minimum per-block corruption probability for block-diffusion training +block_diffusion_logit_alignment: 'same_position' # Supported alignments: same_position, shifted +block_diffusion_canvas_policy: 'all_masked' # Supported canvases: all_masked, seed_and_mask jax_distributed_initialization_timeout: 300 # This is the default timeout in https://github.com/jax-ml/jax/blob/main/jax/_src/distributed.py # Note there are two separate initializations - the jax coordination service (aka jax.distributed.initialize) and the backend (e.g. PjRT), the timeout above refers @@ -1352,4 +1357,3 @@ elastic_backup_kind: "snapshot" elastic_timeout_seconds: 300 elastic_max_retries: 10 elastic_min_slice_count: -1 - diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 769a548422..c9a8e0a6c0 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1646,6 +1646,28 @@ class Distillation(BaseModel): class TrainingLoop(BaseModel): """Configuration for the main training loop, evaluation, and reproducibility.""" + training_objective: Literal["causal_lm", "block_diffusion"] = Field( + "causal_lm", + description="The token-prediction objective used to prepare targets and compute loss.", + ) + block_diffusion_mask_id: int = Field( + -1, + description="The tokenizer mask-token id required by the block-diffusion training objective.", + ) + block_diffusion_min_noise: float = Field( + 1.0e-3, + gt=0.0, + le=1.0, + description="The minimum corruption probability sampled independently for each block.", + ) + block_diffusion_logit_alignment: Literal["same_position", "shifted"] = Field( + "same_position", + description="How model logits align to clean target-token positions.", + ) + block_diffusion_canvas_policy: Literal["all_masked", "seed_and_mask"] = Field( + "all_masked", + description="Whether every block is fully maskable or begins with a clean anchor token.", + ) steps: int = Field( 150_001, ge=-1, @@ -3521,6 +3543,41 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "Block-diffusion attention with attention='autoselected' or attention='flash' requires hardware='tpu'; " "use attention='dot_product' on other hardware." ) + if self.training_objective == "block_diffusion": + if self.attention_type != AttentionType.BLOCK_DIFFUSION.value: + raise ValueError("`training_objective='block_diffusion'` requires `attention_type='block_diffusion'`.") + if self.block_diffusion_mask_id < 0 or self.block_diffusion_mask_id >= self.vocab_size: + raise ValueError( + f"`block_diffusion_mask_id` ({self.block_diffusion_mask_id}) must satisfy " + f"0 <= block_diffusion_mask_id < vocab_size ({self.vocab_size})." + ) + # Block-diffusion attention validation above rejects packing first. + if self.packing: # pragma: no cover + raise ValueError("`training_objective='block_diffusion'` requires `packing=False`.") + if self.mtp_num_layers > 0: + raise ValueError("`training_objective='block_diffusion'` is not compatible with MTP.") + if self.num_vocab_tiling > 1: + raise ValueError("`training_objective='block_diffusion'` is not compatible with vocabulary tiling.") + if self.dataset_type != "hf": + raise ValueError("`training_objective='block_diffusion'` currently requires `dataset_type='hf'`.") + if self.use_dpo: + raise ValueError("`training_objective='block_diffusion'` is not compatible with DPO.") + if self.use_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 = { + ("same_position", "all_masked"), + ("shifted", "seed_and_mask"), + } + model_contract = (self.block_diffusion_logit_alignment, self.block_diffusion_canvas_policy) + if model_contract not in valid_model_contracts: + raise ValueError( + "Block-diffusion training supports only `same_position/all_masked` or `shifted/seed_and_mask`; " + f"received `{model_contract[0]}/{model_contract[1]}`." + ) + if self.block_diffusion_canvas_policy == "seed_and_mask" and self.causal_block_size < 2: + raise ValueError("`block_diffusion_canvas_policy='seed_and_mask'` requires `causal_block_size >= 2`.") if self.quantize_kvcache and not self.kv_quant_axis: raise ValueError("`kv_quant_axis` cannot be empty when quantize_kvcache is True.") if ( diff --git a/src/maxtext/input_pipeline/hf_data_processing.py b/src/maxtext/input_pipeline/hf_data_processing.py index fe3fc44b5e..b8482dc1b8 100644 --- a/src/maxtext/input_pipeline/hf_data_processing.py +++ b/src/maxtext/input_pipeline/hf_data_processing.py @@ -42,6 +42,57 @@ def _get_pad_id(tokenizer): return pad_id +def _get_training_objective_transform( + config: ml_collections.ConfigDict, + *, + shift: bool, + use_dpo: bool, + use_sft: 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. + + 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. + 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. + + Returns: + The objective-specific Grain transform, or None when target shifting is disabled. + + Raises: + ValueError: If the objective is unsupported or block diffusion is combined with + an incompatible post-training or packing mode. + """ + 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.") + if packing: + raise ValueError("Block-diffusion pre-training requires packing=False.") + return input_pipeline_utils.BlockDiffusionCorruption( + block_size=config.causal_block_size, + mask_id=config.block_diffusion_mask_id, + min_noise=config.block_diffusion_min_noise, + logit_alignment=config.block_diffusion_logit_alignment, + canvas_policy=config.block_diffusion_canvas_policy, + axis=1, + ) + if objective != "causal_lm": + raise ValueError(f"Unsupported training objective: {objective}") + if shift and not use_dpo: + return input_pipeline_utils.ShiftData(ignored_ids=[pad_id, bos_token_id], axis=1) + return None + + def vision_sft_preprocessing_pipeline( dataset, config, @@ -354,11 +405,20 @@ def preprocessing_pipeline( max_prompt_length = config.dpo.max_prompt_length operations.append(dpo_utils.DPODataFormatting(pad_id, max_target_length, data_column_names, max_prompt_length)) else: - operations.append(input_pipeline_utils.PadOrTrimToMaxLength(max_target_length, pad_id)) + operations.append(input_pipeline_utils.PadOrTrimToMaxLength(max_target_length, pad_id, config=config)) operations.append(grain.Batch(batch_size=batch_size, drop_remainder=drop_remainder)) - if shift and not use_dpo: - operations.append(input_pipeline_utils.ShiftData(ignored_ids=[pad_id, tokenizer.bos_token_id], axis=1)) + target_transform = _get_training_objective_transform( + config, + shift=shift, + use_dpo=use_dpo, + use_sft=use_sft, + packing=packing, + pad_id=pad_id, + bos_token_id=tokenizer.bos_token_id, + ) + if target_transform is not None: + operations.append(target_transform) # Since HuggingFace IterableDataset does not support access through index # Indexes generated by dummy_index_sampler is not used. diff --git a/src/maxtext/input_pipeline/input_pipeline_utils.py b/src/maxtext/input_pipeline/input_pipeline_utils.py index 45b13b24a5..e2df084794 100644 --- a/src/maxtext/input_pipeline/input_pipeline_utils.py +++ b/src/maxtext/input_pipeline/input_pipeline_utils.py @@ -27,6 +27,7 @@ import numpy as np from grain._src.python.dataset.sources.tfrecord_dataset import _TFRecordReader, _TFRecordDatasetIterator # pylint: disable=protected-access from grain.experimental import TFRecordIterDataset +from maxtext.diffusion.block_diffusion import corruption as block_diffusion_corruption from maxtext.input_pipeline.protos import example_pb2 from maxtext.input_pipeline import tokenizer from maxtext.multimodal import processor as mm_processor @@ -849,14 +850,22 @@ def map( ) -> dict[str, np.ndarray | mm_utils.PreprocessorOutput]: """map to each element""" data_columns = list(element.keys()) + preserve_pad_valued_tokens = ( + 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 isinstance(element[data_column], mm_utils.PreprocessorOutput): raise TypeError("Only 'images' column can be of type PreprocessorOutput.") - element[f"{data_column}_segmentation"] = ( - element[data_column] != self.pad_id # pyrefly: ignore[unsupported-operation] - ) # pyrefly: ignore[unsupported-operation] + if preserve_pad_valued_tokens: + element[f"{data_column}_segmentation"] = np.ones( + element[data_column].shape[0], dtype=np.int32 # pyrefly: ignore[missing-attribute] + ) + else: + element[f"{data_column}_segmentation"] = ( + element[data_column] != self.pad_id # pyrefly: ignore[unsupported-operation] + ) # pyrefly: ignore[unsupported-operation] # pyrefly: ignore[missing-attribute] element[f"{data_column}_segmentation"] = element[ f"{data_column}_segmentation" @@ -878,6 +887,8 @@ 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")): + element[key] = self._pad_text(element[key], self.max_length, 0) # pyrefly: ignore[bad-argument-type] elif "true_length" not in key: element[key] = self._pad_text(element[key], self.max_length, self.pad_id) # pyrefly: ignore[bad-argument-type] return element @@ -1001,6 +1012,55 @@ def map(self, element): return shift_and_refine(element, ignored_ids=self.ignored_ids, axis=self.axis) +@dataclasses.dataclass +class BlockDiffusionCorruption(grain.RandomMapTransform): + """Adapts block-diffusion corruption to the Grain batch contract.""" + + def __init__( + self, + block_size: int, + mask_id: int, + min_noise: float = 1.0e-3, + logit_alignment: str = "same_position", + canvas_policy: str = "all_masked", + axis: int = 1, + ): + self.block_size = block_size + self.mask_id = mask_id + self.min_noise = min_noise + self.logit_alignment = logit_alignment + self.canvas_policy = canvas_policy + self.axis = axis + + def random_map(self, element, rng: np.random.Generator): + """Corrupts inputs while preserving clean targets and input metadata.""" + inputs = np.asarray(element["inputs"]) + targets = np.asarray(element["targets"]) + targets_segmentation = np.asarray(element["targets_segmentation"]) + if inputs.shape != targets.shape or inputs.shape != targets_segmentation.shape: + raise ValueError( + "inputs, targets, and targets_segmentation must have identical shapes, got " + f"{inputs.shape}, {targets.shape}, and {targets_segmentation.shape}" + ) + result = block_diffusion_corruption.corrupt_tokens( + inputs, + targets_segmentation != 0, + rng, + block_size=self.block_size, + mask_id=self.mask_id, + min_noise=self.min_noise, + logit_alignment=self.logit_alignment, + canvas_policy=self.canvas_policy, + axis=self.axis, + ) + output = dict(element) + output["inputs"] = result.inputs + output["targets"] = targets + output["corruption_mask"] = result.corruption_mask.astype(targets_segmentation.dtype) + output["targets_loss_mask"] = result.targets_loss_mask.astype(targets_segmentation.dtype) + return output + + @dataclasses.dataclass class ComputeQwen3OmniPositions(grain.MapTransform): """Computes 3D position IDs for Qwen3-Omni multimodal sequences. diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index 0906c3e786..93a93b401b 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -45,6 +45,7 @@ from flax.nnx import variablelib from maxtext.configs import pyconfig +from maxtext.diffusion.block_diffusion import target_alignment as block_diffusion_target_alignment from maxtext.utils.globals import EPS from maxtext.utils import elastic_utils # Placeholder: internal @@ -107,6 +108,22 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr loss: average loss aux: a dictionary including intermediate_outputs, xent_sum, and total_weights """ + is_block_diffusion = getattr(config, "training_objective", "causal_lm") == "block_diffusion" + if getattr(config, "attention_type", "global") == "block_diffusion" and not is_block_diffusion: + raise ValueError( + "Block-diffusion attention requires target-aligned block-diffusion losses; " + "causal next-token labels would leak within a bidirectional block." + ) + if is_block_diffusion: + required_masks = {"corruption_mask", "targets_loss_mask"} + missing_masks = required_masks - data.keys() + if missing_masks: + raise ValueError(f"Block-diffusion loss requires explicit batch masks; missing {sorted(missing_masks)}") + target_shape = data["targets"].shape + for mask_name in required_masks: + if data[mask_name].shape != target_shape: + raise ValueError(f"{mask_name} must match targets shape; got {data[mask_name].shape} and {target_shape}") + # decimate proportion of data when per_device_batch_size<1 if is_train: for k, v in data.items(): @@ -114,6 +131,9 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr else: for k, v in data.items(): data[k] = v[: config.micro_batch_size_to_eval_on, :] + if is_block_diffusion: + targets_loss_mask = (data["targets_loss_mask"] != 0) & (data["targets_segmentation"] != 0) + target_positions = data.get("targets_position", data["inputs_position"]) mutable_collections = ["intermediates"] if config.mtp_num_layers > 0 and is_train: # The single model.apply call now triggers the entire chain if MTP is enabled: @@ -165,6 +185,13 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0] xent_sum, total_z_loss = vocab_tiling_linen_loss(hidden_states, data, config, model, params, is_train) else: + if is_block_diffusion: + logits = block_diffusion_target_alignment.align_logits_to_targets( + logits, + config.block_diffusion_logit_alignment, + target_positions, + data["targets_segmentation"] != 0, + ) one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size) xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier) @@ -183,9 +210,12 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr debug_sharding=config.debug_sharding, ) - # Mask out paddings at the end of each example. - xent = xent * (data["targets_segmentation"] != 0) - z_loss = z_loss * (data["targets_segmentation"] != 0) + if is_block_diffusion: + xent = xent * targets_loss_mask + z_loss = z_loss * targets_loss_mask + else: + xent = xent * (data["targets_segmentation"] != 0) + z_loss = z_loss * (data["targets_segmentation"] != 0) xent_sum = jnp.sum(xent) total_z_loss = jnp.sum(z_loss) @@ -228,6 +258,13 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr hidden_states = maxtext_utils.get_nested_value(intermediate_outputs, hidden_state_key)[0] xent_sum, total_z_loss = vocab_tiling_nnx_loss(model, hidden_states, data, config, is_train) else: + if is_block_diffusion: + logits = block_diffusion_target_alignment.align_logits_to_targets( + logits, + config.block_diffusion_logit_alignment, + target_positions, + data["targets_segmentation"] != 0, + ) one_hot_targets = jax.nn.one_hot(data["targets"], config.vocab_size) xent, z_loss = max_utils.cross_entropy_with_logits(logits, one_hot_targets, z_loss=config.z_loss_multiplier) @@ -246,14 +283,20 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr debug_sharding=config.debug_sharding, ) - # Mask out paddings at the end of each example. - xent = xent * (data["targets_segmentation"] != 0) - z_loss = z_loss * (data["targets_segmentation"] != 0) + if is_block_diffusion: + xent = xent * targets_loss_mask + z_loss = z_loss * targets_loss_mask + else: + xent = xent * (data["targets_segmentation"] != 0) + z_loss = z_loss * (data["targets_segmentation"] != 0) xent_sum = jnp.sum(xent) total_z_loss = jnp.sum(z_loss) - total_weights = jnp.sum(data["targets_segmentation"] != 0) + if is_block_diffusion: + total_weights = jnp.sum(targets_loss_mask) + else: + total_weights = jnp.sum(data["targets_segmentation"] != 0) # If gradient accumulation is enabled, we don't need to divide xent_sum # by total_weights and then multiply the computed gradient by total_weights, # since it's equivalent to computing the gradient from xent_sum. diff --git a/src/maxtext/utils/gradient_accumulation.py b/src/maxtext/utils/gradient_accumulation.py index 106baee07d..8afd5a22e9 100644 --- a/src/maxtext/utils/gradient_accumulation.py +++ b/src/maxtext/utils/gradient_accumulation.py @@ -162,12 +162,20 @@ def reshape_to_microbatch_accumulations(batch_arr): grad_and_loss, aux = jax.lax.scan( accumulate_gradient, init_grad_and_loss, data, length=config.gradient_accumulation_steps ) + is_block_diffusion = getattr(config, "training_objective", "causal_lm") == "block_diffusion" + if is_block_diffusion: + has_weights = grad_and_loss["total_weights"] > 0 + denominator = jnp.maximum(grad_and_loss["total_weights"], 1) + else: + denominator = grad_and_loss["total_weights"] loss = ( - grad_and_loss["loss"] / grad_and_loss["total_weights"] + grad_and_loss["loss"] / denominator + grad_and_loss["moe_lb_loss"] / config.gradient_accumulation_steps + grad_and_loss["indexer_loss"] / config.gradient_accumulation_steps + grad_and_loss["mtp_loss"] / config.gradient_accumulation_steps ) + if is_block_diffusion: + loss = jnp.where(has_weights, loss, 0.0) raw_grads = grad_and_loss["grad"] if data_parallel_active: # Mark the gradients unreduced over the "data" axis now that we're outside the @@ -177,11 +185,15 @@ def reshape_to_microbatch_accumulations(batch_arr): raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, unreduced_shardings) raw_grads = jax.tree.map(_maybe_shard_with_name, raw_grads, params_shardings) divisor = ( - config.gradient_accumulation_steps - if getattr(config, "use_tunix_gradient_accumulation", False) - else grad_and_loss["total_weights"] + config.gradient_accumulation_steps if getattr(config, "use_tunix_gradient_accumulation", False) else denominator ) - raw_grads = jax.tree_util.tree_map(lambda arr: arr / divisor, raw_grads) + if is_block_diffusion: + raw_grads = jax.tree_util.tree_map( + lambda arr: jnp.where(has_weights, arr / divisor, jnp.zeros_like(arr)), + raw_grads, + ) + else: + raw_grads = jax.tree_util.tree_map(lambda arr: arr / divisor, raw_grads) aux = jax.tree.map(lambda x: jnp.sum(x, axis=0), aux) # pytype: disable=module-attr if is_nnx: diff --git a/src/maxtext/utils/max_utils.py b/src/maxtext/utils/max_utils.py index e6aebce059..a2feea2d1d 100644 --- a/src/maxtext/utils/max_utils.py +++ b/src/maxtext/utils/max_utils.py @@ -957,6 +957,8 @@ def reorder_causal_load_balanced(batch, cp_size, reorder_strategy, hardware="tpu "targets_position", "inputs_segmentation", "targets_segmentation", + "corruption_mask", + "targets_loss_mask", } if hardware in ("gpu", "gpu_multiprocess"): diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 769e86b4bd..c53a03de5c 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -161,6 +161,9 @@ def get_shaped_batch(config, batch_sharding=None): shaped_batch["targets"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) shaped_batch["targets_position"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) shaped_batch["targets_segmentation"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) + if getattr(config, "training_objective", "causal_lm") == "block_diffusion": + shaped_batch["corruption_mask"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) + shaped_batch["targets_loss_mask"] = jax.ShapeDtypeStruct(batch_shape, jnp.int32, sharding=batch_sharding) if config.use_multimodal: image_shape = mm_processor.get_dummy_image_shape_for_init( config.model_name, batch_size=config.micro_batch_size_to_train_on diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index 5a9c0ff922..3a813d438e 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -352,6 +352,108 @@ def test_default_attention_remains_global(self): config = pyconfig.initialize(["", _BASE_CONFIG_PATH, "run_name=test", "steps=1"]) self.assertEqual(config.attention_type, "global") + self.assertEqual(config.training_objective, "causal_lm") + self.assertEqual(config.block_diffusion_mask_id, -1) + + def test_block_diffusion_pretraining_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=7", + "block_diffusion_mask_id=100", + "block_diffusion_min_noise=0.05", + "vocab_size=256", + "max_target_length=2048", + "packing=False", + "dataset_type=hf", + "hf_path=parquet", + "hardware=cpu", + ] + ) + + self.assertEqual(config.training_objective, "block_diffusion") + self.assertEqual(config.block_diffusion_mask_id, 100) + self.assertEqual(config.block_diffusion_min_noise, 0.05) + self.assertEqual(config.block_diffusion_logit_alignment, "same_position") + self.assertEqual(config.block_diffusion_canvas_policy, "all_masked") + + def test_block_diffusion_pretraining_rejects_incompatible_config(self): + base_overrides = { + "run_name": "test", + "steps": 1, + "training_objective": "block_diffusion", + "attention": "dot_product", + "attention_type": "block_diffusion", + "causal_block_size": 32, + "block_diffusion_mask_id": 100, + "block_diffusion_min_noise": 0.05, + "vocab_size": 256, + "max_target_length": 2048, + "packing": False, + "dataset_type": "hf", + "hf_path": "parquet", + "hardware": "cpu", + } + cases = ( + ({"attention_type": "global"}, "attention_type='block_diffusion'"), + ({"block_diffusion_mask_id": -1}, "block_diffusion_mask_id"), + ({"block_diffusion_mask_id": 256}, "block_diffusion_mask_id"), + ({"block_diffusion_min_noise": 0.0}, "block_diffusion_min_noise"), + ({"packing": True}, "packing=False"), + ({"mtp_num_layers": 1}, "MTP"), + ({"num_vocab_tiling": 2}, "vocabulary tiling"), + ({"dataset_type": "grain"}, "dataset_type='hf'"), + ({"use_dpo": True}, "DPO"), + ({"use_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"), + ( + { + "causal_block_size": 1, + "block_diffusion_logit_alignment": "shifted", + "block_diffusion_canvas_policy": "seed_and_mask", + }, + "causal_block_size >= 2", + ), + ) + for overrides, expected_regex in cases: + with self.subTest(overrides=overrides): + values = base_overrides | overrides + argv = ["", _BASE_CONFIG_PATH, *(f"{key}={value}" for key, value in values.items())] + with self.assertRaisesRegex((ValueError, pydantic.ValidationError), expected_regex): + pyconfig.initialize(argv) + + def test_shifted_block_diffusion_requires_seeded_canvas(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", + "block_diffusion_logit_alignment=shifted", + "block_diffusion_canvas_policy=seed_and_mask", + "vocab_size=256", + "packing=False", + "dataset_type=hf", + "hf_path=parquet", + "hardware=cpu", + ] + ) + + self.assertEqual(config.block_diffusion_logit_alignment, "shifted") + self.assertEqual(config.block_diffusion_canvas_policy, "seed_and_mask") @unittest.mock.patch.dict(os.environ, {pyconfig.yaml_key_to_env_key("steps"): "123"}) def test_env_override(self): diff --git a/tests/unit/gradient_accumulation_nnx_test.py b/tests/unit/gradient_accumulation_nnx_test.py index 478fe9a746..567fdc79cc 100644 --- a/tests/unit/gradient_accumulation_nnx_test.py +++ b/tests/unit/gradient_accumulation_nnx_test.py @@ -14,6 +14,8 @@ """Unit tests for the NNX branch of gradient_accumulation_loss_and_grad.""" +# pylint: disable=too-many-positional-arguments + import unittest from dataclasses import dataclass @@ -34,6 +36,8 @@ class _Cfg: shard_mode: int = ShardMode.AUTO ici_data_parallelism: int = 1 debug_sharding: bool = False + training_objective: str = "causal_lm" + use_tunix_gradient_accumulation: bool = False class _TinyNNX(nnx.Module): @@ -64,7 +68,27 @@ def _fake_loss_fn(model, config, data, dropout_rng, params, is_train=True): "indexer_loss": jnp.array(0.0), "mtp_loss": jnp.array(0.0), } - return xent_sum / total_weights, aux + return xent_sum, aux + + +def _zero_weight_loss_fn(model, config, data, dropout_rng, params, is_train=True): + """Produces model-dependent sums with a zero token denominator.""" + del config, dropout_rng, params, is_train + xent_sum = jnp.sum(model(data["inputs"]) ** 2) + 1.0 + aux = { + "xent_sum": xent_sum, + "total_weights": jnp.array(0.0), + "moe_lb_loss": jnp.array(0.0), + "indexer_loss": jnp.array(0.0), + "mtp_loss": jnp.array(0.0), + } + return xent_sum, aux + + +def _normalized_loss_fn(model, config, data, dropout_rng, params, is_train=True): + """Produces the per-microbatch normalized loss used by Tunix accumulation.""" + xent_sum, aux = _fake_loss_fn(model, config, data, dropout_rng, params, is_train=is_train) + return xent_sum / aux["total_weights"], aux class TestGradientAccumulationNNX(unittest.TestCase): @@ -114,6 +138,55 @@ def test_nnx_path_runs_and_returns_grad_for_every_param(self): for g in grad_leaves: self.assertTrue(jnp.all(jnp.isfinite(g))) + def test_positive_weights_match_direct_causal_batch(self): + """The guarded denominator preserves ordinary causal accumulation.""" + kernel = self.model.linear.kernel.get_value() + bias = self.model.linear.bias.get_value() + + def direct_loss(kernel, bias): + predictions = self.data["inputs"] @ kernel + bias + return jnp.mean((predictions - self.data["targets"]) ** 2) + + expected_loss, expected_grads = jax.value_and_grad(direct_loss, argnums=(0, 1))(kernel, bias) + loss, _, raw_grads = gradient_accumulation.gradient_accumulation_loss_and_grad( + _fake_loss_fn, + self.cfg, + self.model, + params=None, + params_shardings=self._params_shardings(), + data=self.data, + dropout_rng=None, + ) + + np.testing.assert_allclose(loss, expected_loss, rtol=1e-6) + np.testing.assert_allclose(raw_grads["linear"]["kernel"].get_value(), expected_grads[0], rtol=1e-6) + np.testing.assert_allclose(raw_grads["linear"]["bias"].get_value(), expected_grads[1], rtol=1e-6) + + def test_tunix_block_diffusion_uses_accumulation_step_divisor(self): + self.cfg.training_objective = "block_diffusion" + self.cfg.use_tunix_gradient_accumulation = True + kernel = self.model.linear.kernel.get_value() + bias = self.model.linear.bias.get_value() + + def direct_loss(kernel, bias): + predictions = self.data["inputs"] @ kernel + bias + return jnp.mean((predictions - self.data["targets"]) ** 2) + + expected_loss, expected_grads = jax.value_and_grad(direct_loss, argnums=(0, 1))(kernel, bias) + loss, _, raw_grads = gradient_accumulation.gradient_accumulation_loss_and_grad( + _normalized_loss_fn, + self.cfg, + self.model, + params=None, + params_shardings=self._params_shardings(), + data=self.data, + dropout_rng=None, + ) + + np.testing.assert_allclose(loss, expected_loss, rtol=1e-6) + np.testing.assert_allclose(raw_grads["linear"]["kernel"].get_value(), expected_grads[0], rtol=1e-6) + np.testing.assert_allclose(raw_grads["linear"]["bias"].get_value(), expected_grads[1], rtol=1e-6) + def test_nnx_path_updates_model_rest_state_after_scan(self): """After accumulation, nnx.update is called on the model with the rest_state from the scan. @@ -151,6 +224,26 @@ def test_nnx_with_shard_optimizer_over_data_casts_to_bf16(self): for g in jax.tree.leaves(raw_grads): self.assertTrue(jnp.all(jnp.isfinite(g))) + def test_zero_total_weights_returns_zero_loss_and_gradients(self): + self.cfg.training_objective = "block_diffusion" + for use_tunix_gradient_accumulation in (False, True): + with self.subTest(use_tunix_gradient_accumulation=use_tunix_gradient_accumulation): + self.cfg.use_tunix_gradient_accumulation = use_tunix_gradient_accumulation + loss, _, raw_grads = gradient_accumulation.gradient_accumulation_loss_and_grad( + _zero_weight_loss_fn, + self.cfg, + self.model, + params=None, + params_shardings=self._params_shardings(), + data=self.data, + dropout_rng=None, + ) + + self.assertEqual(float(loss), 0.0) + for gradient in jax.tree.leaves(raw_grads): + self.assertTrue(jnp.all(jnp.isfinite(gradient))) + self.assertTrue(jnp.all(gradient == 0)) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/hf_data_processing_test.py b/tests/unit/hf_data_processing_test.py index 262c56ff9b..3b2c7bf738 100644 --- a/tests/unit/hf_data_processing_test.py +++ b/tests/unit/hf_data_processing_test.py @@ -15,7 +15,9 @@ """Tests for Hugging Face data processing.""" import sys +from types import SimpleNamespace import unittest +from unittest import mock import os.path import jax @@ -25,6 +27,7 @@ from maxtext.configs import pyconfig from maxtext.input_pipeline import hf_data_processing from maxtext.input_pipeline import input_pipeline_interface +from maxtext.input_pipeline import input_pipeline_utils from maxtext.common.gcloud_stub import is_decoupled from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory @@ -124,5 +127,161 @@ def get_first_batch(iterator): self.assertTrue((train_batch1["targets"] == train_batch2["targets"]).all()) # pytype: disable=unsupported-operands +class TrainingObjectiveTransformTest(unittest.TestCase): + """Tests the pre-training objective boundary in the HF pipeline.""" + + def _block_diffusion_config(self): + return SimpleNamespace( + elastic_enabled=False, + training_objective="block_diffusion", + causal_block_size=4, + block_diffusion_mask_id=99, + block_diffusion_min_noise=0.05, + block_diffusion_canvas_policy="seed_and_mask", + block_diffusion_logit_alignment="shifted", + ) + + def _pipeline_operations(self, config, *, shift): + """Builds the lightweight non-packing pipeline and returns its operations.""" + dataset = mock.MagicMock() + dataset.select_columns.return_value = dataset + tokenizer = SimpleNamespace(pad_token_id=0, unk_token_id=1, bos_token_id=2) + data_source = [object()] + dataloader = object() + iterator = object() + mesh = SimpleNamespace(size=1) + + with ( + mock.patch.object(hf_data_processing.transformers.AutoTokenizer, "from_pretrained", return_value=tokenizer), + mock.patch.object(hf_data_processing.input_pipeline_utils, "HFDataSource", return_value=data_source), + mock.patch.object(hf_data_processing.grain, "DataLoader", return_value=dataloader) as data_loader, + mock.patch.object( + hf_data_processing.multihost_dataloading, + "MultiHostDataLoadIterator", + return_value=iterator, + ), + ): + result = hf_data_processing.preprocessing_pipeline( + dataloading_host_index=0, + dataloading_host_count=1, + global_mesh=mesh, + dataset=dataset, + config=config, + data_column_names=("text",), + tokenize=False, + tokenizer_path="unused", + hf_access_token=None, + global_batch_size=1, + max_target_length=8, + shuffle=False, + data_shuffle_seed=0, + packing=False, + shift=shift, + use_dpo=False, + use_sft=False, + ) + + self.assertIs(result, iterator) + return data_loader.call_args.kwargs["operations"] + + def test_default_objective_keeps_next_token_shift(self): + transform = hf_data_processing._get_training_objective_transform( # pylint: disable=protected-access + SimpleNamespace(), + shift=True, + use_dpo=False, + use_sft=False, + packing=False, + pad_id=0, + bos_token_id=1, + ) + + self.assertIsInstance(transform, input_pipeline_utils.ShiftData) + self.assertEqual(transform.ignored_ids, [0, 1]) + + def test_causal_objective_without_shift_has_no_transform(self): + transform = hf_data_processing._get_training_objective_transform( # pylint: disable=protected-access + SimpleNamespace(training_objective="causal_lm"), + shift=False, + use_dpo=False, + use_sft=False, + packing=False, + pad_id=0, + bos_token_id=1, + ) + + self.assertIsNone(transform) + + def test_unknown_objective_is_rejected(self): + with self.assertRaisesRegex(ValueError, "Unsupported training objective"): + hf_data_processing._get_training_objective_transform( # pylint: disable=protected-access + SimpleNamespace(training_objective="unknown"), + shift=False, + use_dpo=False, + use_sft=False, + packing=False, + pad_id=0, + bos_token_id=1, + ) + + def test_block_diffusion_replaces_next_token_shift(self): + transform = hf_data_processing._get_training_objective_transform( # pylint: disable=protected-access + self._block_diffusion_config(), + shift=True, + use_dpo=False, + use_sft=False, + packing=False, + pad_id=0, + bos_token_id=1, + ) + + self.assertIsInstance(transform, input_pipeline_utils.BlockDiffusionCorruption) + self.assertEqual(transform.block_size, 4) + self.assertEqual(transform.mask_id, 99) + self.assertEqual(transform.min_noise, 0.05) + self.assertEqual(transform.logit_alignment, "shifted") + self.assertEqual(transform.canvas_policy, "seed_and_mask") + + def test_preprocessing_pipeline_installs_block_diffusion_transform(self): + operations = self._pipeline_operations(self._block_diffusion_config(), shift=True) + + self.assertTrue(any(isinstance(operation, input_pipeline_utils.PadOrTrimToMaxLength) for operation in operations)) + self.assertIsInstance(operations[-1], input_pipeline_utils.BlockDiffusionCorruption) + + def test_preprocessing_pipeline_omits_disabled_causal_shift(self): + operations = self._pipeline_operations( + SimpleNamespace(elastic_enabled=False, training_objective="causal_lm"), + shift=False, + ) + + self.assertTrue(any(isinstance(operation, input_pipeline_utils.PadOrTrimToMaxLength) for operation in operations)) + self.assertFalse( + any( + isinstance(operation, (input_pipeline_utils.BlockDiffusionCorruption, input_pipeline_utils.ShiftData)) + for operation in operations + ) + ) + + def test_block_diffusion_rejects_packing_and_post_training_modes(self): + base_args = { + "shift": True, + "use_dpo": False, + "use_sft": 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: + with self.subTest(overrides=overrides), self.assertRaisesRegex(ValueError, expected_message): + hf_data_processing._get_training_objective_transform( # pylint: disable=protected-access + self._block_diffusion_config(), + **(base_args | overrides), + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/input_pipeline_utils_test.py b/tests/unit/input_pipeline_utils_test.py index c64dba591a..889cb5db11 100644 --- a/tests/unit/input_pipeline_utils_test.py +++ b/tests/unit/input_pipeline_utils_test.py @@ -15,8 +15,84 @@ """Unit tests for input_pipeline_utils.""" import unittest - -from maxtext.input_pipeline.input_pipeline_utils import compute_file_sharding +from types import SimpleNamespace + +import numpy as np + +from maxtext.input_pipeline.input_pipeline_utils import BlockDiffusionCorruption, compute_file_sharding, PadOrTrimToMaxLength + + +class BlockDiffusionPaddingTest(unittest.TestCase): + """Checks that tokenizer padding never becomes diffusion supervision.""" + + def test_nonzero_token_pad_id_keeps_metadata_padding_zero(self): + clean = PadOrTrimToMaxLength( + max_length=6, + pad_id=7, + config=SimpleNamespace(training_objective="block_diffusion"), + ).map( + { + "inputs": np.asarray([11, 12, 13], dtype=np.int32), + "targets": np.asarray([11, 12, 13], dtype=np.int32), + } + ) + corrupted = BlockDiffusionCorruption(block_size=4, mask_id=99, min_noise=1.0, axis=0).random_map( + clean, np.random.default_rng(0) + ) + + padding = np.arange(6) >= 3 + np.testing.assert_array_equal(clean["inputs"][padding], 7) + np.testing.assert_array_equal(clean["inputs_segmentation"][padding], 0) + np.testing.assert_array_equal(clean["targets_segmentation"][padding], 0) + np.testing.assert_array_equal(clean["inputs_position"][padding], 0) + np.testing.assert_array_equal(clean["targets_position"][padding], 0) + np.testing.assert_array_equal(corrupted["corruption_mask"][padding], 0) + np.testing.assert_array_equal(corrupted["targets_loss_mask"][padding], 0) + np.testing.assert_array_equal(corrupted["inputs"][padding], 7) + + def test_pad_valued_source_token_remains_valid(self): + clean = PadOrTrimToMaxLength( + max_length=4, + pad_id=7, + config=SimpleNamespace(training_objective="block_diffusion"), + ).map( + { + "inputs": np.asarray([11, 7], dtype=np.int32), + "targets": np.asarray([11, 7], dtype=np.int32), + } + ) + + np.testing.assert_array_equal(clean["inputs"], [11, 7, 7, 7]) + np.testing.assert_array_equal(clean["inputs_segmentation"], [1, 1, 0, 0]) + np.testing.assert_array_equal(clean["targets_segmentation"], [1, 1, 0, 0]) + + def test_causal_padding_preserves_legacy_metadata_pad_value(self): + clean = PadOrTrimToMaxLength( + max_length=4, + pad_id=7, + config=SimpleNamespace(training_objective="causal_lm"), + ).map( + { + "inputs": np.asarray([11, 7], dtype=np.int32), + "targets": np.asarray([11, 7], dtype=np.int32), + } + ) + + np.testing.assert_array_equal(clean["inputs_segmentation"], [1, 0, 7, 7]) + np.testing.assert_array_equal(clean["targets_segmentation"], [1, 0, 7, 7]) + np.testing.assert_array_equal(clean["inputs_position"], [0, 1, 7, 7]) + np.testing.assert_array_equal(clean["targets_position"], [0, 1, 7, 7]) + + def test_corruption_rejects_mismatched_batch_shapes(self): + with self.assertRaisesRegex(ValueError, "must have identical shapes"): + BlockDiffusionCorruption(block_size=4, mask_id=99).random_map( + { + "inputs": np.asarray([11, 12], dtype=np.int32), + "targets": np.asarray([11, 12, 13], dtype=np.int32), + "targets_segmentation": np.ones(2, dtype=np.int32), + }, + np.random.default_rng(0), + ) class ComputeFileShardingNormalCaseTest(unittest.TestCase): diff --git a/tests/unit/max_utils_test.py b/tests/unit/max_utils_test.py index 097d85d491..5587842e65 100644 --- a/tests/unit/max_utils_test.py +++ b/tests/unit/max_utils_test.py @@ -23,10 +23,12 @@ import jax from jax import numpy as jnp from jax import random +from maxtext.common.common_types import ReorderStrategy from maxtext.configs import pyconfig from maxtext.utils import max_utils from maxtext.utils.train_utils import setup_train_loop from tests.utils.test_helpers import get_test_config_path +import numpy as np import optax import pytest @@ -655,6 +657,28 @@ def test_reorder_roundtrip(self): # 3. Assert roundtrip is lossless self.assertTrue(jnp.allclose(x, restored, rtol=1e-5, atol=1e-6)) + def test_block_diffusion_masks_are_reordered_with_tokens(self): + corruption_mask = jnp.arange(16, dtype=jnp.int32).reshape(1, 16) + targets_loss_mask = corruption_mask + 100 + batch = { + "inputs": corruption_mask + 200, + "corruption_mask": corruption_mask, + "targets_loss_mask": targets_loss_mask, + } + + reordered = max_utils.reorder_causal_load_balanced( + batch, + cp_size=2, + reorder_strategy=ReorderStrategy.DUAL_CHUNK_SWAP, + hardware="cpu", + ) + + np.testing.assert_array_equal(reordered["corruption_mask"], max_utils.reorder_sequence(corruption_mask, 2)) + np.testing.assert_array_equal( + reordered["targets_loss_mask"], + max_utils.reorder_sequence(targets_loss_mask, 2), + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py index 472ce303d3..58f9bf00e9 100644 --- a/tests/unit/maxtext_utils_test.py +++ b/tests/unit/maxtext_utils_test.py @@ -1151,8 +1151,16 @@ def test_linen_in_shardings_includes_rng(self): class TestGetShapedBatch(unittest.TestCase): """Tests for get_shaped_batch.""" - def _make_cfg(self, *, enable_diloco=False, use_multimodal=False, use_audio=False, model_name="llama3.1-8b"): - """Build a minimal config mock for get_shaped_batch tests.""" + def _make_cfg( + self, + *, + enable_diloco=False, + use_multimodal=False, + use_audio=False, + model_name="llama3.1-8b", + training_objective="causal_lm", + ): + """Builds the config subset consumed by get_shaped_batch.""" cfg = MagicMock() cfg.enable_diloco = enable_diloco cfg.global_batch_size_to_load = 4 @@ -1160,6 +1168,7 @@ def _make_cfg(self, *, enable_diloco=False, use_multimodal=False, use_audio=Fals cfg.use_multimodal = use_multimodal cfg.use_audio = use_audio cfg.model_name = model_name + cfg.training_objective = training_objective if enable_diloco: cfg.num_diloco_replicas = 2 return cfg @@ -1182,6 +1191,21 @@ def test_standard_shape(self): expected_shape = (cfg.global_batch_size_to_load, cfg.max_target_length) self.assertEqual(batch["inputs"].shape, expected_shape) + def test_block_diffusion_masks_are_in_shaped_batch(self): + cfg = self._make_cfg(training_objective="block_diffusion") + + batch = maxtext_utils.get_shaped_batch(cfg) + + self.assertEqual(batch["corruption_mask"].shape, batch["inputs"].shape) + self.assertEqual(batch["targets_loss_mask"].shape, batch["inputs"].shape) + self.assertEqual(batch["targets_loss_mask"].dtype, jnp.int32) + + def test_causal_shaped_batch_has_no_diffusion_masks(self): + batch = maxtext_utils.get_shaped_batch(self._make_cfg()) + + self.assertNotIn("corruption_mask", batch) + self.assertNotIn("targets_loss_mask", batch) + def test_diloco_shape(self): cfg = self._make_cfg(enable_diloco=True) batch = maxtext_utils.get_shaped_batch(cfg) diff --git a/tests/unit/pre_train_loss_mask_test.py b/tests/unit/pre_train_loss_mask_test.py new file mode 100644 index 0000000000..c58349ac98 --- /dev/null +++ b/tests/unit/pre_train_loss_mask_test.py @@ -0,0 +1,251 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests explicit target-loss masking in the pre-training loss.""" + +# pylint: disable=too-many-positional-arguments + +from dataclasses import dataclass +import unittest +from unittest import mock + +from flax import linen as nn +from flax import nnx +import jax +import jax.numpy as jnp + +from maxtext.trainers.pre_train import train as pre_train + + +@dataclass +class _Config: + """Configuration subset consumed by pre-training loss_fn.""" + + micro_batch_size_to_train_on: int = 2 + micro_batch_size_to_eval_on: int = 2 + vocab_size: int = 8 + z_loss_multiplier: float = 0.0 + enable_dropout: bool = False + use_multimodal: bool = False + use_indexer: bool = False + indexer_sparse_training: bool = False + indexer_loss_scaling_factor: float = 0.0 + num_vocab_tiling: int = 1 + num_experts: int = 1 + routed_bias: bool = False + routed_bias_update_rate: float = 0.0 + mtp_num_layers: int = 0 + mtp_eval_target_module: int = 0 + use_qk_clip: bool = False + use_tunix_gradient_accumulation: bool = False + gradient_accumulation_steps: int = 1 + shard_mode: int = 0 + debug_sharding: bool = False + weight_sparsity_n: int = 0 + weight_sparsity_m: int = 0 + attention_type: str = "global" + training_objective: str = "causal_lm" + block_diffusion_logit_alignment: str = "same_position" + block_diffusion_canvas_policy: str = "all_masked" + causal_block_size: int = 4 + + +class _UniformNnxDecoder(nnx.Module): + """Returns uniform logits through the NNX call contract.""" + + def __init__(self, vocab_size): + self.vocab_size = vocab_size + self.mesh = jax.make_mesh((1, 1, 1, 1), ("data", "fsdp", "expert", "context")) + + def __call__( + self, + decoder_input_tokens, + decoder_positions, + decoder_segment_ids=None, + encoder_images=None, + encoder_image_masks=None, + enable_dropout=False, + decoder_target_tokens=None, + decoder_target_mask=None, + ): + del decoder_positions, decoder_segment_ids, encoder_images, encoder_image_masks + del enable_dropout, decoder_target_tokens, decoder_target_mask + return jnp.zeros((*decoder_input_tokens.shape, self.vocab_size), dtype=jnp.float32) + + +class _UniformLinenDecoder(nn.Module): + """Returns uniform logits through the Linen call contract.""" + + vocab_size: int + mesh: object + + @nn.compact + def __call__( + self, + decoder_input_tokens, + decoder_positions, + decoder_segment_ids=None, + encoder_images=None, + encoder_image_masks=None, + enable_dropout=False, + decoder_target_tokens=None, + decoder_target_mask=None, + ): + del decoder_positions, decoder_segment_ids, encoder_images, encoder_image_masks + del enable_dropout, decoder_target_tokens, decoder_target_mask + return jnp.zeros((*decoder_input_tokens.shape, self.vocab_size), dtype=jnp.float32) + + +def _make_data(include_loss_mask=True): + """Builds a batch whose explicit loss mask differs from segmentation.""" + data = { + "inputs": jnp.zeros((2, 4), dtype=jnp.int32), + "inputs_position": jnp.broadcast_to(jnp.arange(4), (2, 4)), + "inputs_segmentation": jnp.ones((2, 4), dtype=jnp.int32), + "targets": jnp.zeros((2, 4), dtype=jnp.int32), + "targets_segmentation": jnp.asarray([[1, 1, 1, 1], [1, 1, 1, 0]], dtype=jnp.int32), + "corruption_mask": jnp.asarray([[1, 0, 1, 0], [0, 1, 0, 0]], dtype=jnp.int32), + } + if include_loss_mask: + data["targets_loss_mask"] = jnp.asarray([[1, 0, 1, 0], [0, 1, 0, 0]], dtype=jnp.int32) + return data + + +class PreTrainLossMaskTest(unittest.TestCase): + """Checks explicit masks in both Linen and NNX loss branches.""" + + def setUp(self): + super().setUp() + self.config = _Config() + self.per_token_xent = jnp.arange(1, 9, dtype=jnp.float32).reshape(2, 4) + self.per_token_z_loss = self.per_token_xent / 10.0 + + def _cross_entropy_patch(self): + return mock.patch.object( + pre_train.max_utils, + "cross_entropy_with_logits", + return_value=(self.per_token_xent, self.per_token_z_loss), + ) + + def _use_block_diffusion(self): + self.config.attention_type = "block_diffusion" + self.config.training_objective = "block_diffusion" + + def _linen_model_and_variables(self, data): + """Initializes the test Linen decoder for the supplied batch.""" + mesh = jax.make_mesh((1, 1, 1, 1), ("data", "fsdp", "expert", "context")) + model = _UniformLinenDecoder(vocab_size=self.config.vocab_size, mesh=mesh) + variables = model.init( + jax.random.key(0), + data["inputs"], + data["inputs_position"], + decoder_segment_ids=data["inputs_segmentation"], + decoder_target_tokens=data["targets"], + decoder_target_mask=data["targets_segmentation"], + ) + return model, variables + + def _assert_explicit_mask_result(self, loss, aux): + expected_mask = _make_data()["targets_loss_mask"] != 0 + expected_xent = jnp.sum(self.per_token_xent * expected_mask) + expected_z_loss = jnp.sum(self.per_token_z_loss * expected_mask) / 3.0 + self.assertEqual(int(aux["total_weights"]), 3) + self.assertAlmostEqual(float(aux["xent_sum"]), float(expected_xent)) + self.assertAlmostEqual(float(loss), float(expected_xent / 3.0)) + self.assertAlmostEqual(float(aux["z_loss"]), float(expected_z_loss)) + + def test_nnx_loss_uses_targets_loss_mask(self): + self._use_block_diffusion() + model = _UniformNnxDecoder(self.config.vocab_size) + with self._cross_entropy_patch(): + loss, aux = pre_train.loss_fn(model, self.config, _make_data(), None, None, is_train=True) + + self._assert_explicit_mask_result(loss, aux) + + def test_linen_loss_uses_targets_loss_mask(self): + self._use_block_diffusion() + data = _make_data() + model, variables = self._linen_model_and_variables(data) + with self._cross_entropy_patch(): + loss, aux = pre_train.loss_fn(model, self.config, data, jax.random.key(1), variables, is_train=True) + + self._assert_explicit_mask_result(loss, aux) + + def test_linen_causal_loss_skips_diffusion_alignment(self): + data = _make_data(include_loss_mask=False) + model, variables = self._linen_model_and_variables(data) + with self._cross_entropy_patch(): + loss, aux = pre_train.loss_fn(model, self.config, data, jax.random.key(1), variables, is_train=True) + + expected_mask = data["targets_segmentation"] != 0 + expected_xent = jnp.sum(self.per_token_xent * expected_mask) + self.assertEqual(int(aux["total_weights"]), 7) + self.assertAlmostEqual(float(aux["xent_sum"]), float(expected_xent)) + self.assertAlmostEqual(float(loss), float(expected_xent / 7.0)) + + def test_causal_fallback_uses_targets_segmentation(self): + data = _make_data(include_loss_mask=False) + model = _UniformNnxDecoder(self.config.vocab_size) + with self._cross_entropy_patch(): + loss, aux = pre_train.loss_fn(model, self.config, data, None, None, is_train=True) + + expected_mask = data["targets_segmentation"] != 0 + expected_xent = jnp.sum(self.per_token_xent * expected_mask) + self.assertEqual(int(aux["total_weights"]), 7) + self.assertAlmostEqual(float(aux["xent_sum"]), float(expected_xent)) + self.assertAlmostEqual(float(loss), float(expected_xent / 7.0)) + + def test_zero_explicit_mask_has_finite_zero_loss(self): + self._use_block_diffusion() + data = _make_data() + data["targets_loss_mask"] = jnp.zeros_like(data["targets_loss_mask"]) + model = _UniformNnxDecoder(self.config.vocab_size) + with self._cross_entropy_patch(): + loss, aux = pre_train.loss_fn(model, self.config, data, None, None, is_train=True) + + self.assertEqual(int(aux["total_weights"]), 0) + self.assertEqual(float(aux["xent_sum"]), 0.0) + self.assertTrue(bool(jnp.isfinite(loss))) + self.assertEqual(float(loss), 0.0) + + def test_causal_loss_rejects_block_bidirectional_attention(self): + self.config.attention_type = "block_diffusion" + model = _UniformNnxDecoder(self.config.vocab_size) + + with self.assertRaisesRegex(ValueError, "would leak"): + pre_train.loss_fn(model, self.config, _make_data(), None, None, is_train=True) + + def test_block_diffusion_requires_all_explicit_masks(self): + self._use_block_diffusion() + model = _UniformNnxDecoder(self.config.vocab_size) + + for mask_name in ("corruption_mask", "targets_loss_mask"): + with self.subTest(mask_name=mask_name): + data = _make_data() + del data[mask_name] + with self.assertRaisesRegex(ValueError, mask_name): + pre_train.loss_fn(model, self.config, data, None, None, is_train=True) + + def test_block_diffusion_rejects_mismatched_mask_shape(self): + self._use_block_diffusion() + data = _make_data() + data["targets_loss_mask"] = data["targets_loss_mask"][:, :-1] + model = _UniformNnxDecoder(self.config.vocab_size) + + with self.assertRaisesRegex(ValueError, "must match targets shape"): + pre_train.loss_fn(model, self.config, data, None, None, is_train=True) + + +if __name__ == "__main__": + unittest.main()