PaLM-rlhf-pytorch: A PyTorch Skeleton for RLHF, Not a Chatbot

What PaLM-rlhf-pytorch actually ships

The repository is a training scaffold for reinforcement learning from human feedback, built on a PaLM-style autoregressive transformer. The README describes it as an implementation of RLHF on top of the PaLM architecture, and the project description frames it as ChatGPT but with PaLM. That framing is about the pipeline shape, not about capability.

The audience is narrow. This is for engineers who understand PPO, already have a pretrained transformer checkpoint, and want the RLHF stages expressed as PyTorch modules they can read and modify. The README answers the obvious question directly: there is no trained model. It calls the repository the ship and the overall map, and says millions of dollars of compute plus data are still needed to reach the right point in parameter space. Nothing here generates useful answers out of the box.

That is worth stating before anything else, because the package name and the ChatGPT comparison invite the opposite assumption. The top-level layout is small: a palm_rlhf_pytorch/ package, an examples.py, a train.py, a data/ directory, a LICENSE and a pyproject.toml.

The three-stage pipeline the README walks through

The mechanism is the standard RLHF decomposition, and the README presents it as three scripts in sequence.

First, a PaLM model is trained as an ordinary autoregressive transformer. The README's example constructs PaLM with num_tokens, dim, depth and an optional flash_attn flag, runs a forward pass with return_loss = True, calls backward, and notes that after much training you can call generate.

Second, a RewardModel wraps that transformer with causal = False and a num_binned_output parameter, described as a rating scale such as 1 to 5. The forward call takes a sequence, a prompt_mask marking which tokens are prompt and which are response, and labels. The README notes that in the original paper the reward model could not be finetuned from a pretrained transformer without overfitting, but that LoRA finetuning is offered anyway because it remains open research. That is an honest caveat about a stage many write-ups treat as settled.

Third, RLHFTrainer takes the PaLM model, the reward model and prompt_token_ids, and runs trainer.train(num_episodes = 50000). Generation after training is a sampling step: trainer.generate(2048, prompt = prompts[0], num_samples = 10) draws ten candidates and the reward model picks the best. The TODO list confirms the actor and critic are separate, with the critic cloned from the base transformer and given its own LoRA, and that non-LoRA finetuning is also supported.

Installing palm-rlhf-pytorch and running a first reward model

The README gives one install line. The package is on PyPI as palm-rlhf-pytorch, while the import name is palm_rlhf_pytorch with underscores.

pip install palm-rlhf-pytorch

pyproject.toml requires Python 3.6 or later and pulls in torch>=2.2 along with accelerate, beartype, einops>=0.8.0, einx>=0.3.0, memmap-replay-buffer>=0.1.4, hl-gauss-pytorch>=0.1.19, torch-einops-utils, adam-atan2-pytorch, x-mlps-pytorch and tqdm. Note that the version in pyproject.toml is 0.7.5 while the most recent release listed is 0.7.1, so pin deliberately if reproducibility matters to you.

The smallest real exercise is the reward model stage, which does not require a trained policy. The README's example builds a non-causal PaLM, wraps it, and feeds a mock batch:

import torch
from palm_rlhf_pytorch import PaLM, RewardModel

palm = PaLM(num_tokens = 20000, dim = 512, depth = 12, causal = False)
reward_model = RewardModel(palm, num_binned_output = 5).cuda()

seq = torch.randint(0, 20000, (1, 1024)).cuda()
prompt_mask = torch.zeros(1, 1024).bool().cuda()
labels = torch.randint(0, 5, (1,)).cuda()

loss = reward_model(seq, prompt_mask = prompt_mask, labels = labels)
loss.backward()

If that runs, the wiring is correct: a scalar loss comes back and gradients flow into the wrapped transformer. It says nothing about whether the reward model is any good, which depends entirely on the human feedback data you supply. The README does not document a data format for that feedback, so expect to read the source before feeding real annotations.

Where the scaffold stops being useful

The largest limitation is stated in the README itself: no trained model, no data, no compute budget. Anyone expecting a downloadable checkpoint has misread the project. The reward model needs curated human feedback, and the README gives no pipeline for collecting it; a Django and htmx interface for that is still an unchecked box in the TODO list.

Several other items in that list are unfinished and touch core behaviour. Sampling with variable-length prompts is not done, with the README arguing the bottleneck is human feedback rather than sampling. Writing PPO memory to a memmapped numpy file is not done. Finetuning only the penultimate N layers of the actor or critic is not done. Instrumentation with Hugging Face accelerate and wandb is not done. A sentiment network as a stand-in reward model has not been tested.

There is also a design question the README raises against itself. It calls Direct Preference Optimization a potential successor and says all the code in the repository becomes roughly a binary cross entropy loss, under five lines. That is the author arguing that the reward model plus PPO machinery here may be more apparatus than the problem needs. If your goal is preference tuning rather than studying PPO, the three-stage structure is a cost, not a feature.

One more boundary: this is not a serving stack. There is no inference server, no API, no tokenizer, and no chat template anywhere in the repository. It is a training library.

How it compares with TRLX and Open-Assistant

The README points to two other efforts rather than pretending to be alone in the space. CarperAI had been building an RLHF framework for large language models, trlx, for months before ChatGPT shipped, and Yannic Kilcher is working on an open sourced implementation under LAION-AI/Open-Assistant.

The difference in approach is scope. trlx is a framework aimed at running RLHF across large models with distributed training concerns handled for you. PaLM-rlhf-pytorch is a single-author implementation that keeps the model classes small and readable: PaLM, RewardModel, RLHFTrainer, each constructible in a few lines. If you want to modify how the critic is initialized or how rewards are normalized, this codebase is easier to hold in your head. If you want to fine-tune a 7B model on multiple nodes next week, a framework built for that is the better fit.

Open-Assistant is a project rather than a library, oriented toward producing an open chat model with collected data. It answers a different question: what do we train, rather than how does the loop work.

The README also credits Enrico for integrating Flash Attention from PyTorch 2.0, which is why flash_attn appears as a constructor flag on PaLM.

Maintenance, licensing and upgrade exposure

The repository is not archived, and the last push was on 2026-07-27. Releases listed are 0.7.1 on 2025-09-19, 0.6.3 and 0.6.2 both on 2025-09-02. pyproject.toml declares version 0.7.5, which is ahead of the newest listed release, so the published package and the repository state can drift apart between releases. Pin an exact version in your own requirements file rather than relying on the latest tag to match main.

The dependency list is broad for a research scaffold: accelerate, beartype, einops, einx, memmap-replay-buffer, hl-gauss-pytorch, torch-einops-utils, adam-atan2-pytorch, x-mlps-pytorch, tqdm, and torch at 2.2 or newer. Several are also lucidrains packages, which means upgrades to those libraries can ripple into this one. Expect to spend time on dependency resolution before you spend time on training.

The licence is MIT, declared both in the LICENSE file and in the pyproject.toml license field. MIT is permissive, so commercial use and modification are permitted subject to the licence text. That covers this repository only. The PaLM paper and the RLHF literature it cites are separate works, and the human feedback data you train on carries its own terms. Nothing here is legal advice; read the LICENSE and your data agreements.

Editorial conclusion

Adopt this if you already have a pretrained PaLM checkpoint, a reward model and a cluster, and you want the three-stage RLHF loop expressed as readable PyTorch modules rather than a service. Do not adopt it if you want a chatbot today; the README states there is no trained model and that reaching one needs millions of dollars of compute and data, and the author's own note points at Direct Preference Optimization as a potential successor that collapses most of this code into a binary cross entropy loss. Before committing, verify that the version in pyproject.toml, 0.7.5, matches the release you intend to pin, and read palm_rlhf_pytorch/ and train.py to confirm the actor, critic and reward model classes expose the constructor arguments your checkpoint expects.

Frequently asked questions

Official sources

1. Issues
2. License: MIT
3. lucidrains/PaLM-rlhf-pytorch on GitHub
4. README
5. Releases