PaLM-rlhf-pytorch: the RLHF pipeline as a skeleton, not a model

What PaLM-rlhf-pytorch actually provides

The repository is an implementation of reinforcement learning from human feedback on top of the PaLM architecture, described in its own README as work in progress. It is not a product and not a checkpoint. The FAQ answers the obvious question first: there is no trained model, and the README compares the code to a ship and a map, with the compute and data needed to reach the destination still missing. So the deliverable is the pipeline: an autoregressive transformer, a reward model that scores sequences against binned human ratings, and a trainer that runs PPO against that reward. The intended reader is someone who wants to read or modify the RLHF loop in PyTorch rather than call an API. The README also points at Laion's Discord for people trying to replicate something like ChatGPT in the open, which sets the ambition level honestly.

Three classes, three training stages

The README walks through the stages in order. First, PaLM is trained like any other autoregressive transformer, with the example constructing it from num_tokens, dim, depth and an optional flash_attn flag, then calling it with return_loss=True and backpropagating. Second, RewardModel wraps a PaLM instance built with causal=False and a num_binned_output argument, described in the example as a rating from 1 to 5. The forward call takes the sequence, a boolean prompt_mask marking which positions are prompt and which are response, and integer labels, and returns a loss. Third, RLHFTrainer receives the trained palm, the trained reward_model, and a tensor of prompt token ids, then runs trainer.train(num_episodes=50000). Generation afterwards is not plain sampling: trainer.generate(2048, prompt=prompts[0], num_samples=10) produces ten candidates and the reward model picks the best. That last step is the clearest signal of what the package is for, since it treats the reward model as both a training signal and a reranker.

The critic is a LoRA clone, and that is a design choice

The Todo list records a completed item: clone base transformer with separate lora for critic. A second completed item adds the alternative, also allow for non-LoRA based finetuning. This is the part of the repository that carries the most practical weight. PPO needs a value estimate alongside the policy, and duplicating a full PaLM for that doubles the memory cost of the largest object in the system. A LoRA-adapted copy keeps the critic small relative to the actor. The cost is that the critic's capacity is bounded by the adapter, which may or may not be enough depending on how far the policy moves from its initialization. The README does not quantify that trade-off, and it does not state which path is the default. Anyone reading the code should check before assuming the cheaper option is active. The same list notes that the reward model can be finetuned with LoRA even though the original paper could not finetune a reward model from a pretrained transformer without overfitting; the author keeps the option open and labels it open research.

Getting it running: install and the minimum loop

Installation is a single command, pip install palm-rlhf-pytorch. There is no server, no CLI and no config file; everything is Python objects. The README's usage section gives the sequence explicitly. Construct PaLM with num_tokens, dim, depth and optionally flash_attn=True, then call palm(seq, return_loss=True) and backward. Construct RewardModel(palm, num_binned_output=5) over a PaLM built with causal=False, then call reward_model(seq, prompt_mask=prompt_mask, labels=labels). Construct RLHFTrainer(palm=..., reward_model=..., prompt_token_ids=prompts) and call trainer.train(num_episodes=50000). The README shows palm.load('./path/to/pretrained/palm.pt') and reward_model.load('./path/to/pretrained/reward_model.pt'), which means checkpoint persistence is your responsibility; the package does not define a format beyond whatever those load calls expect. Note also that the README's own example reuses the same palm object when building the reward model, which is a documentation shortcut rather than a recommended recipe. In a real run the reward model needs its own weights.

Where the pipeline breaks down

The most concrete limitation is stated by the project itself: no trained model, and the compute to produce one is out of reach for most readers. That is not a defect in the code, but it does mean the repository cannot be evaluated end to end by anyone without a cluster. Several Todo items are still unchecked and each one marks a rough edge: memory in PPO is not yet written to a memmapped numpy file, sampling with variable-length prompts is not working, and finetuning only the penultimate N layers of the actor or critic is not implemented. The variable-length item matters more than it looks, because the README's own example feeds a fixed prompt tensor of shape (50000, 512), and real prompts are not uniform. There is also no Hugging Face accelerate integration and no wandb instrumentation yet, so distributed training and experiment tracking are on you. Finally, the README itself points at Direct Preference Optimization as a potential successor and notes that under DPO all the code in the repository reduces to a binary cross entropy loss in under five lines. That is the author's own assessment of how much of this machinery a newer method might replace.

Alternatives and the actual difference

The README names CarperAI's trlx as an RLHF framework for large language models that predates the ChatGPT release, and LAION's Open-Assistant as another open implementation. The difference is mostly scope and shape. trlx is described in the README as a framework, which implies a training harness with configuration and scaling concerns handled at the framework level; PaLM-rlhf-pytorch is a set of nn.Module classes you compose yourself, with no config layer and no opinion about distributed execution. If you want to read PPO, the reward model and the actor-critic wiring in one file each and change them, the single-author package is easier to hold in your head. If you want to run a training job across many GPUs with logging and checkpoint management handled for you, the framework is the better starting point, and the README's own Todo list (accelerate, wandb) shows which direction this repository has not yet gone. The DPO paper is a third path and not a drop-in replacement: it removes the reward model and the PPO loop rather than reimplementing them, which is why the README says the code here collapses into a binary cross entropy loss.

Licence, releases and what maintenance looks like

The repository is MIT licensed, which permits commercial use, modification and redistribution provided the copyright notice and permission notice are retained. That is a permissive default and it matters here, because a company building on this code inherits no copyleft obligation from the licence itself. It says nothing about the licences of PaLM weights, preference datasets or pretrained reward models you bring in, and those are separate questions. On maintenance, the release history shows 0.6.2 and 0.6.3 in September 2025 and 0.7.1 later that month, with the last push in July 2026, so the package is being touched rather than abandoned. Upgrade cost is low in the ordinary sense, since the public surface is a handful of classes and the README documents their arguments. The real cost is not versioning; it is that the unfinished Todo items are the parts that stand between the skeleton and a working run, and each release may or may not close one of them. Reading the Todo list before pinning a version tells you more about readiness than the version number does.

Editorial conclusion

Adopt it if you want a readable PyTorch reference for the three-stage RLHF loop and you already have a pretrained PaLM checkpoint and human preference data of your own. Do not adopt it expecting a usable assistant: the FAQ states plainly that there is no trained model and that reaching one costs millions of dollars of compute plus data. Before committing, verify three things in the source tree: that RLHFTrainer still constructs the critic as a LoRA clone of the base transformer, that the reward model path you intend to use (LoRA finetuning or full finetuning) is the one the current release supports, and that the variable-length prompt sampling item on the Todo list is still open, because it is listed as unfinished.

Official sources

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