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()