[maxtext] Add block-diffusion CFT and SFT

src/dependencies/extra_deps/post_train_github_deps.txt

@@ -1,3 +1,3 @@
-google-tunix @ https://github.com/google/tunix/archive/c4ec573d29e4c3a3955b348256d464b119c8a6d1.zip
+google-tunix @ https://github.com/ethannnnnn/tunix/archive/0ab812a8908e971d3f61d92844f07e12b6ee4806.zip
 tpu-inference @ https://github.com/vllm-project/tpu-inference/archive/7ecc401e6faefe3f793f3e4a8e6e2dbc49eca868.zip
 vllm @ git+https://github.com/vllm-project/vllm@0ba2aa35a81dcc3246b26291368b53fa2389c7d7

src/maxtext/configs/base.yml

@@ -1032,6 +1032,7 @@ decode_sampling_temperature: 1.
 eval_start_step: 0 # start eval when train step is >= eval_start_step
 eval_interval: -1  # the specific number of train step between eval_step
 eval_steps: -1  # run this number of steps for eval, recommend setting this to prevent error due to running out of evel data
+loss_is_preaveraged: false # Whether eval hooks receive a loss already averaged across eval batches
 target_eval_loss: 0.  # early stop once reaching target eval_loss
 abort_on_nan_loss: true # Check for NaN and abort if found in training loss
 abort_on_inf_loss: true # Check for Inf and abort if found in training loss

src/maxtext/configs/types.py

@@ -1719,6 +1719,10 @@ class TrainingLoop(BaseModel):
       -1,
       description="Number of steps to run for each evaluation. -1 runs on entire eval split.",
   )
+  loss_is_preaveraged: bool = Field(
+      False,
+      description="Whether the evaluation hook receives a loss already averaged across evaluation batches.",
+  )
   target_eval_loss: float = Field(
       0.0,
       description="If set, training will stop early when this evaluation loss is reached.",
@@ -3727,8 +3731,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
         raise ValueError("`training_objective='block_diffusion'` currently requires `dataset_type='hf'`.")
       if self.use_dpo:
         raise ValueError("`training_objective='block_diffusion'` is not compatible with DPO.")
-      if self.use_sft:
-        raise ValueError("`training_objective='block_diffusion'` currently supports pre-training only.")
       if self.use_multimodal or self.use_audio:
         raise ValueError("`training_objective='block_diffusion'` currently supports text-only training.")
       valid_model_contracts = {

src/maxtext/input_pipeline/hf_data_processing.py

@@ -48,17 +48,19 @@ def _get_training_objective_transform(
     shift: bool,
     use_dpo: bool,
     use_sft: bool,
+    completion_only: bool,
     packing: bool,
     pad_id: int,
     bos_token_id: int | None,
 ) -> input_pipeline_utils.ShiftData | input_pipeline_utils.BlockDiffusionCorruption | None:
-  """Selects target preparation for causal or block-diffusion pre-training.
+  """Selects target preparation for causal or block-diffusion training.
 
   Args:
     config: Training configuration containing the objective-specific settings.
     shift: Whether causal language-model targets should be shifted by one token.
     use_dpo: Whether the pipeline is preparing direct-preference data.
     use_sft: Whether the pipeline is preparing supervised fine-tuning data.
+    completion_only: Whether SFT supervision is restricted to completion tokens.
     packing: Whether multiple examples are packed into each sequence.
     pad_id: Token ID used to pad causal language-model examples.
     bos_token_id: Beginning-of-sequence token ID, or None when unavailable.
@@ -68,23 +70,22 @@ def _get_training_objective_transform(
 
   Raises:
     ValueError: If the objective is unsupported or block diffusion is combined with
-      an incompatible post-training or packing mode.
+      DPO or packing.
   """
   objective = getattr(config, "training_objective", "causal_lm")
   if objective == "block_diffusion":
-    if use_sft:
-      raise ValueError("This block-diffusion integration currently supports pre-training only.")
     if use_dpo:
-      raise ValueError("Block-diffusion pre-training is not compatible with DPO.")
+      raise ValueError("Block-diffusion training is not compatible with DPO.")
     if packing:
-      raise ValueError("Block-diffusion pre-training requires packing=False.")
+      raise ValueError("Block-diffusion training requires packing=False.")
     return input_pipeline_utils.BlockDiffusionCorruption(
         block_size=config.causal_block_size,
         mask_id=config.block_diffusion_mask_id,
         min_noise=config.block_diffusion_min_noise,
         logit_alignment=config.block_diffusion_logit_alignment,
         canvas_policy=config.block_diffusion_canvas_policy,
         axis=1,
+        completion_only=bool(use_sft and completion_only),
     )
   if objective != "causal_lm":
     raise ValueError(f"Unsupported training objective: {objective}")
@@ -389,6 +390,7 @@ def preprocessing_pipeline(
             completion_only=sft_train_on_completion_only,
             max_target_length=max_target_length,
             unk_id=pad_id,
+            training_objective=getattr(config, "training_objective", "causal_lm"),
         )
     )
     data_column_names = ("inputs", "targets")
@@ -424,6 +426,7 @@ def preprocessing_pipeline(
       shift=shift,
       use_dpo=use_dpo,
       use_sft=use_sft,
+      completion_only=sft_train_on_completion_only,
       packing=packing,
       pad_id=pad_id,
       bos_token_id=tokenizer.bos_token_id,

src/maxtext/input_pipeline/input_pipeline_utils.py

@@ -329,33 +329,51 @@ def tokenization(example, hf_tokenizer, truncation, max_length, column_names):
 @dataclasses.dataclass
 class SFTPromptMasking(grain.MapTransform):
   """Construct inputs and targets for SFT training. Concat prompt and completion to generate inputs.
-  For targets, if train on completion only, the prompt will be masked by unk_id. Otherwise the same as inputs.
+  Causal SFT can mask prompt targets with ``unk_id``. Block diffusion keeps
+  clean same-position targets and emits role-derived completion metadata.
   """
 
-  def __init__(self, text_column_name, completion_only, max_target_length, unk_id=0):
+  def __init__(
+      self,
+      text_column_name,
+      completion_only,
+      max_target_length,
+      unk_id=0,
+      training_objective="causal_lm",
+  ):
+    if training_objective not in ("causal_lm", "block_diffusion"):
+      raise ValueError(f"Unsupported training objective: {training_objective}")
     self.text_column_name = text_column_name
     self.completion_only = completion_only
     self.max_target_length = max_target_length
     self.unk_id = unk_id
+    self.training_objective = training_objective
 
   def map(self, element):
     """
     Maps a single dataset element to an SFT training instance.
     It concatenates the prompt and completion to form the `inputs` sequence.
     For the `targets` sequence:
-    - If `self.completion_only` is `True`, the prompt portion of the
-      concatenated sequence is masked using `self.unk_id`.
+    - If `self.completion_only` is `True`, causal SFT masks the prompt portion
+      with `self.unk_id`; block-diffusion SFT keeps clean targets and records
+      completion eligibility separately.
     - If `self.completion_only` is `False`, the target sequence is
       identical to the input sequence.
     """
-    inputs, targets = [], []
+    is_block_diffusion = self.training_objective == "block_diffusion"
+    inputs, targets, completion_mask = [], [], []
     for i, text in enumerate(element[self.text_column_name]):
       inputs += text
-      targets += [self.unk_id] * len(text) if self.completion_only and element["is_prompt"][i] else text
-    return {
+      is_prompt = element["is_prompt"][i]
+      targets += [self.unk_id] * len(text) if self.completion_only and is_prompt and not is_block_diffusion else text
+      completion_mask += [not is_prompt] * len(text)
+    output = {
         "inputs": np.asarray(inputs[: self.max_target_length], dtype=np.int32),
         "targets": np.asarray(targets[: self.max_target_length], dtype=np.int32),
     }
+    if is_block_diffusion:
+      output["completion_mask"] = np.asarray(completion_mask[: self.max_target_length], dtype=np.int32)
+    return output
 
 
 @dataclasses.dataclass
@@ -854,7 +872,7 @@ def map(
         self.config is not None and getattr(self.config, "training_objective", "causal_lm") == "block_diffusion"
     )
     for data_column in data_columns:
-      if data_column != "images":
+      if data_column not in ("images", "completion_mask"):
         if isinstance(element[data_column], mm_utils.PreprocessorOutput):
           raise TypeError("Only 'images' column can be of type PreprocessorOutput.")
 
@@ -887,7 +905,7 @@ def map(
 
         element["images"] = self._pad_image_and_mask(element["images"])  # pyrefly: ignore[bad-argument-type]
 
-      elif preserve_pad_valued_tokens and key.endswith(("_segmentation", "_position")):
+      elif key == "completion_mask" or (preserve_pad_valued_tokens and key.endswith(("_segmentation", "_position"))):
         element[key] = self._pad_text(element[key], self.max_length, 0)  # pyrefly: ignore[bad-argument-type]
       elif "true_length" not in key:
         element[key] = self._pad_text(element[key], self.max_length, self.pad_id)  # pyrefly: ignore[bad-argument-type]
@@ -1026,6 +1044,7 @@ class BlockDiffusionCorruption(grain.RandomMapTransform):
   logit_alignment: str = "same_position"
   canvas_policy: str = "all_masked"
   axis: int = 1
+  completion_only: bool = False
 
   def random_map(self, element, rng: np.random.Generator):
     """Corrupts inputs while preserving clean targets and input metadata."""
@@ -1037,9 +1056,51 @@ def random_map(self, element, rng: np.random.Generator):
           "inputs, targets, and targets_segmentation must have identical shapes, got "
           f"{inputs.shape}, {targets.shape}, and {targets_segmentation.shape}"
       )
+    validity_mask = targets_segmentation != 0
+    raw_completion_mask = element.get("completion_mask")
+    if raw_completion_mask is None:
+      if self.completion_only:
+        if "completion_mask" not in element:
+          raise ValueError("completion-only block diffusion requires an explicit completion_mask")
+        raise ValueError("completion-only block diffusion requires a non-None completion_mask")
+      completion_mask = validity_mask
+    else:
+      completion_mask = np.asarray(raw_completion_mask) != 0
+    if completion_mask.shape != inputs.shape:
+      raise ValueError(f"completion_mask must match inputs shape; received {completion_mask.shape} and {inputs.shape}")
+    if np.any(completion_mask & ~validity_mask):
+      raise ValueError("completion_mask must be a subset of valid target positions")
+    can_validate_block_order = (
+        self.completion_only
+        and self.block_size > 0
+        and inputs.ndim > 0
+        and -inputs.ndim <= self.axis < inputs.ndim
+        and inputs.shape[self.axis % inputs.ndim] > 0
+    )
+    if can_validate_block_order:
+      axis = self.axis % inputs.ndim
+      sequence_length = inputs.shape[axis]
+      block_count = (sequence_length + self.block_size - 1) // self.block_size
+      padded_length = block_count * self.block_size
+
+      def _as_blocks(mask):
+        rows = np.moveaxis(mask, axis, -1).reshape(-1, sequence_length)
+        padded = np.pad(rows, ((0, 0), (0, padded_length - sequence_length)))
+        return padded.reshape(rows.shape[0], block_count, self.block_size)
+
+      validity_blocks = _as_blocks(validity_mask)
+      completion_blocks = _as_blocks(completion_mask)
+      seen_completion = np.maximum.accumulate(completion_blocks, axis=-1)
+      if np.any(seen_completion & validity_blocks & ~completion_blocks):
+        raise ValueError(
+            "completion-only block diffusion cannot place a prompt token after "
+            "a completion token in the same diffusion block; split the example "
+            "or align the next prompt to a block boundary"
+        )
+    supervision_mask = completion_mask if self.completion_only else validity_mask
     result = block_diffusion_corruption.corrupt_tokens(
         inputs,
-        targets_segmentation != 0,
+        supervision_mask,
         rng,
         block_size=self.block_size,
         mask_id=self.mask_id,