Merge pull request #3997 from AI-Hypercomputer:igorts/delete-legacy-dpo

PiperOrigin-RevId: 924994024

Author: Google-ML-Automation

src/maxtext/__init__.py

@@ -37,7 +37,6 @@
 
 from maxtext.configs import pyconfig
 from maxtext.models import models
-from maxtext.trainers.post_train.dpo import dpo_utils
 from maxtext.utils import maxtext_utils
 from maxtext.utils import model_creation_utils
 

src/maxtext/common/metric_logger.py

@@ -225,8 +225,8 @@ def _log_eval_metrics(self, metrics, step):
               f"avg_mtp_acceptance_rate={scalars['eval/avg_mtp_acceptance_rate_percent']:.2f}%",
           ]
       )
-    if self.config.use_dpo:
-      log_parts.append(f"dpo_reward_accuracy={scalars['eval/dpo_reward_accuracy']:.3f}")
+    if "eval/avg_dpo_reward_accuracy" in scalars:
+      log_parts.append(f"dpo_reward_accuracy={scalars['eval/avg_dpo_reward_accuracy']:.3f}")
     max_logging.log(", ".join(log_parts))
 
   def _log_running_eval_metrics(self, metrics, step):
@@ -421,10 +421,6 @@ def _accumulate_eval_metrics(self, metrics):
         scalar.get("evaluation/mtp_acceptance_rate_percent", 0.0)
     )
     self.cumulative_eval_metrics["scalar"]["eval/z_loss"] += float(scalar.get("evaluation/z_loss", 0.0))
-    if self.config.use_dpo:
-      self.cumulative_eval_metrics["scalar"]["eval/dpo_reward_accuracy"] += float(
-          scalar.get("evaluation/dpo_reward_accuracy", 0.0)
-      )
 
   def record_train_metrics(self, metrics, step, step_time):
     """Records training metrics for the current step."""
@@ -454,8 +450,7 @@ def _finalize_eval_metrics(self, train_step):
     cumulative["eval/avg_mtp_loss"] = cumulative["eval/mtp_loss"] / eval_step_count
     cumulative["eval/avg_mtp_acceptance_rate_percent"] = cumulative["eval/mtp_acceptance_rate_percent"] / eval_step_count
     cumulative["eval/avg_z_loss"] = cumulative["eval/z_loss"] / eval_step_count
-    if self.config.use_dpo:
-      cumulative["eval/dpo_reward_accuracy"] = cumulative["eval/dpo_reward_accuracy"] / eval_step_count
+
     self.write_metrics(self.cumulative_eval_metrics, train_step, metric_type="eval")
     self._pending_eval_step_count = 0
     if self.config.target_eval_loss and eval_loss <= self.config.target_eval_loss:

src/maxtext/common/train_state_nnx.py

@@ -30,31 +30,21 @@ class TrainStateNNX(nnx.Module):
     {"params": {...}, "opt_state": {}...}
   TrainStateNNX state pytree:
     {"model": {...}, "optimizer": {"opt_state": {...}}}
-
-  For DPO (Direct Preference Optimization), an optional `reference_model`
-  carries a frozen copy of the same architecture used to compute reference
-  log-probabilities. Only `model` is updated by `apply_gradients`; the
-  reference is held alongside so it is sharded, jit-traced, and checkpointed
-  with the rest of the train state.
   """
 
   def __init__(
       self,
       model: nnx.Module,
       optimizer: nnx.Optimizer | None,
-      reference_model: nnx.Module | None = None,
   ):
     self.model = model
     self.optimizer = optimizer
-    if reference_model is not None:
-      self.reference_model = reference_model
 
   def apply_gradients(self, grads: Any):
     """Mimics the Linen apply_gradients function.
 
     Updates the optimizer state, applies updates to parameters, and increments
-    the step counter. Only updates `self.model`; `self.reference_model` (if
-    present) is left untouched.
+    the step counter. Only updates `self.model`.
     """
     if self.optimizer is None:
       raise RuntimeError(
@@ -88,9 +78,7 @@ def _cast_step(step, dtype):
   values.
   """
   if isinstance(step, jax.ShapeDtypeStruct):
-    return jax.ShapeDtypeStruct(
-        step.shape, dtype, sharding=getattr(step, "sharding", None)
-    )
+    return jax.ShapeDtypeStruct(step.shape, dtype, sharding=getattr(step, "sharding", None))
   return jnp.asarray(step, dtype=dtype)
 
 
@@ -117,10 +105,7 @@ def _wrap_mu_nu_with_params(state):
   """Wraps mu/nu under an inner 'params' key (the Linen collection)."""
   if not isinstance(state, dict):
     return state
-  return {
-      k: {"params": v} if k in ("mu", "nu") and isinstance(v, dict) else v
-      for k, v in state.items()
-  }
+  return {k: {"params": v} if k in ("mu", "nu") and isinstance(v, dict) else v for k, v in state.items()}
 
 
 def _as_chain_index(key):
@@ -172,23 +157,14 @@ def _strip_mu_nu_params(state):
   if not isinstance(state, dict):
     return state
   return {
-      k: (
-          v["params"]
-          if k in ("mu", "nu") and isinstance(v, dict) and "params" in v
-          else v
-      )
-      for k, v in state.items()
+      k: (v["params"] if k in ("mu", "nu") and isinstance(v, dict) and "params" in v else v) for k, v in state.items()
   }
 
 
 def _opt_state_from_linen(opt_state):
   """Inverse of `_opt_state_to_linen`: Linen list-with-None -> NNX int-keyed dict."""
   if isinstance(opt_state, list):
-    return {
-        i: _strip_mu_nu_params(e)
-        for i, e in enumerate(opt_state)
-        if isinstance(e, dict)
-    }
+    return {i: _strip_mu_nu_params(e) for i, e in enumerate(opt_state) if isinstance(e, dict)}
   if not isinstance(opt_state, dict):
     return opt_state
   return {0: _strip_mu_nu_params(opt_state)}

src/maxtext/input_pipeline/grain_data_processing.py

@@ -263,32 +263,6 @@ def pretrain_preprocessing_pipeline(
   return dataset
 
 
-def dpo_preprocessing_pipeline(
-    dataset,
-    config,
-    data_columns,
-    tokenize,
-    grain_worker_count,
-    grain_per_worker_buffer_size,
-):
-  """Use grain to pre-process the dataset and return iterators for dpo fine-tuning"""
-  dataset = data_processing_utils.parse_and_keep_features(dataset, config, data_columns, tokenize)
-  tokenizer_model, pad_id = data_processing_utils.get_tokenizer_and_pad_id(config)
-
-  if tokenize:
-    dataset = dataset.map(grain_tokenizer.TokenizeAndTrim(data_columns, config.max_target_length, tokenizer_model))
-
-  batch_size = config.global_batch_size_to_load // jax.process_count()
-  # DPO scores full sequences, so no shift.
-  dataset = data_processing_utils.format_and_batch(
-      dataset, config, batch_size, pad_id, data_columns, tokenizer_model, shift=False
-  )
-  dataset = data_processing_utils.apply_multiprocessing_and_prefetch(
-      dataset, config, grain_worker_count, grain_per_worker_buffer_size
-  )
-  return dataset
-
-
 def _format_chat_template_grain(element, data_columns, tokenizer_model):
   """Grain-compatible mapping function to format raw columns into conversational messages."""
   # Convert raw columns to conversational messages
@@ -376,8 +350,6 @@ def sft_preprocessing_pipeline(
 
 def _get_pipeline_fn(config):
   """Returns the appropriate preprocessing pipeline function based on config."""
-  if config.use_dpo:
-    return dpo_preprocessing_pipeline
   if config.use_sft:
     return sft_preprocessing_pipeline
   return pretrain_preprocessing_pipeline

src/maxtext/input_pipeline/tfds_data_processing.py

@@ -91,7 +91,6 @@ def preprocessing_pipeline(
     shift: bool = True,
     drop_remainder: bool = True,
     prefetch_size=tf.data.experimental.AUTOTUNE,
-    use_dpo: bool = False,
     hf_access_token: str = "",
 ):
   """pipeline for preprocessing TFDS dataset."""
@@ -115,15 +114,11 @@ def preprocessing_pipeline(
           "Set tokenize_train_data or tokenize_eval_data to True if your dataset needs tokenization."
       )
 
-  if not use_dpo:
-    assert len(data_column_names) == 1
-    dataset = dataset.map(
-        lambda x: input_pipeline_utils.normalize_features(x, data_column_names[0]), num_parallel_calls=AUTOTUNE
-    )
-  else:
-    dataset = dataset.map(lambda x: {col: x[col] for col in data_column_names}, num_parallel_calls=AUTOTUNE)
-
-  data_column_names = data_column_names if use_dpo else ("inputs", "targets")
+  assert len(data_column_names) == 1
+  dataset = dataset.map(
+      lambda x: input_pipeline_utils.normalize_features(x, data_column_names[0]), num_parallel_calls=AUTOTUNE
+  )
+  data_column_names = ("inputs", "targets")
 
   tokenizer_model = input_pipeline_utils.get_tokenizer(tokenizer_path, tokenizer_type, add_bos, add_eos, hf_access_token)
   if tokenizer_model.pad_id is not None:
@@ -144,7 +139,7 @@ def preprocessing_pipeline(
   if max_target_length > 0:
     # in pre-training we can take upto max_length+1 because there would be truncation by
     # 1 token for both inputs and targets
-    extra_tokens = 1 if not use_dpo else 0
+    extra_tokens = 1
     dataset = dataset.map(
         lambda x: input_pipeline_utils.truncate_to_max_allowable_length(x, max_target_length + extra_tokens),
         num_parallel_calls=AUTOTUNE,
@@ -157,13 +152,13 @@ def preprocessing_pipeline(
   dataset = dataset.repeat(num_epochs)
 
   # Shift inputs for teacher-forced training
-  if shift and not use_dpo:
+  if shift:
     dataset = dataset.map(
         input_pipeline_utils.shift_data_by_truncation, num_parallel_calls=tf.data.AUTOTUNE, deterministic=True
     )
 
   # Perform greedy sequence packing and batching
-  if pack_examples and not use_dpo:
+  if pack_examples:
     dataset = sequence_packing.pack_dataset(dataset, max_target_length, pad_id)
     dataset = dataset.batch(global_batch_size // jax.process_count(), drop_remainder=drop_remainder)
   else:
@@ -223,7 +218,6 @@ def make_tfds_train_iterator(
         add_eos=config.add_eos,
         num_epochs=config.num_epoch,
         pack_examples=config.packing,
-        use_dpo=config.use_dpo,
         hf_access_token=config.hf_access_token,
     )
     return multihost_dataloading.MultiHostDataLoadIterator(
@@ -248,7 +242,6 @@ def make_tfds_train_iterator(
         add_eos=config.add_eos,
         num_epochs=config.num_epoch,
         pack_examples=config.packing,
-        use_dpo=config.use_dpo,
         hf_access_token=config.hf_access_token,
     )
     global_shape = (config.global_batch_size_to_load, config.max_target_length)
@@ -289,7 +282,6 @@ def make_tfds_eval_iterator(
         add_bos=config.add_bos,
         add_eos=config.add_eos,
         pack_examples=config.packing,
-        use_dpo=config.use_dpo,
         hf_access_token=config.hf_access_token,
     )
     return multihost_dataloading.MultiHostDataLoadIterator(
@@ -317,7 +309,6 @@ def make_tfds_eval_iterator(
         add_bos=config.add_bos,
         add_eos=config.add_eos,
         pack_examples=config.packing,
-        use_dpo=config.use_dpo,
         hf_access_token=config.hf_access_token,
     )
     global_shape = (config.global_batch_size_to_load_eval, config.max_target_length)