Reproducing OpenAI’s RLHF results has been something of a white whale for researchers without access to very specific, very expensive hardware. A new project from Costa Huang’s team changes that, offering a minimal reference implementation that matches OpenAI’s learning curves on stylistic tasks like sentiment and descriptiveness almost exactly. The codebase, built in PyTorch, is designed for educational use, and the authors have published all training logs publicly on Weights & Biases.
Getting the original OpenAI code to run at all required some creative problem-solving. The original TensorFlow 1.x implementation only works on 8x V100 GPUs with 32GB of memory each—a configuration essentially only available through AWS’s p3dn.24xlarge instances. It won’t run on a single V100 because there’s no gradient accumulation, and it’s incompatible with A100s due to CUDA version issues. The original book dataset was also lost during OpenAI’s GCP-to-Azure migration, so the team substituted Hugging Face’s bookcorpusdataset.
The technical details matter enormously here. One subtle but critical finding: the reward model and policy’s value head don’t just look at the model’s response. They concatenate the query and response together and process the full sequence, producing a reward for every token based on its preceding context. That’s a meaningful departure from how many people assume RLHF works, and getting it wrong would silently degrade performance.
Perhaps the most interesting discovery is a numerical issue in PyTorch’s Adam optimizer that doesn’t exist in TensorFlow. The difference causes aggressive updates during training, which could explain why some PyTorch RLHF implementations have struggled to match OpenAI’s results. This is exactly the kind of implementation detail that doesn’t show up in papers but determines whether your model actually learns. For practitioners building on this work, checking optimizer behavior should be step one before debugging anything else.
