Function _split_dpo_state

src/MaxText/dpo_utils.py:25–29

Split DPO state to separate reference parameters.

(state)

def _split_dpo_state(state):
  """Split DPO state to separate reference parameters."""
  reference_params = state.params["reference_params"]
  new_state = state.replace(params={k: v for k, v in state.params.items() if k != "reference_params"})
  return new_state, reference_params


def dpo_loss_fn(model, config, data, dropout_rng, params, reference_params, is_train=True):