RLHF take-home aligning a small open model
A take-home over a few days: align a small open model with preference data and a parameter-efficient adapter on a single GPU, push the result to the Hugging Face Hub, and write up what happened, including what failed and why.
Choose the model, the data and the method yourself and justify them. The review covers your evaluation, your failure analysis and your reasoning about the loss.
As a warm-up, fill in the core of a DPO step without a framework:
"""The core of the DPO step, without a framework, so it can be unit-tested.
Fill in the three functions; the training script calls them per batch."""
def sequence_logp(token_logps, loss_mask):
"""Sum of log-probabilities over the response tokens of one sequence.
loss_mask is 1 on response tokens and 0 on prompt and padding tokens."""
raise NotImplementedError
def dpo_loss_and_grad(policy_chosen, policy_rejected, ref_chosen, ref_rejected, beta=0.1):
"""DPO loss for one preference pair from sequence log-probabilities.
Returns (loss, d_loss/d_policy_chosen, d_loss/d_policy_rejected).
Must stay finite for very large margins."""
raise NotImplementedError
def dpo_metrics(pairs, beta=0.1):
"""pairs: [(policy_chosen, policy_rejected, ref_chosen, ref_rejected)].
Returns {"loss", "accuracy", "margin", "chosen_reward", "rejected_reward"}
averaged over pairs; a reward is beta * (policy - ref) and accuracy is the
share of pairs whose chosen reward beats the rejected one."""
raise NotImplementedError
Hugging FaceTake-homeSign in to see the worked answer
The full answer, chapter by chapter, with what interviewers expect at each level. Free with an account: one full answer per company.