Advanced model alignment: RLHF, RLAIF, and RLVR in Dataiku#

The process of creating Large Language Models (LLMs) involves several distinct stages. While the pre-training phase creates a model that can predict the next token, it is the alignment phase that truly refines how it generates tokens to better suit its role.

While Supervised Fine-Tuning (SFT) is usually used to teach a model to follow instructions, alignment ensures the model’s behavior matches human preferences, follows specific reasoning paths, or adheres to objective truth. These techniques transition a general next-token predictor into a helpful assistant.

In this tutorial, we will explore advanced model alignment in Dataiku, with three primary methods used in the industry today: Reinforcement Learning from Human Feedback (RLHF), AI Feedback (RLAIF), and Verifiable Rewards (RLVR).

Prerequisites#

To follow along with the code examples in this tutorial, you will need:

  • Dataiku >= 14.1

  • Python >= 3.10

  • A code environment with containerized GPU support (container runtime addition GPU support for Torch 2) and the following packages:

    torch
transformers
datasets
trl
peft
bitsandbytes
accelerate

  • An LLM Mesh connection to HuggingFace.

Reinforcement Learning from Human Feedback (RLHF)#

RLHF is the industry standard for aligning models to human preferences. It uses preference-based optimization to train a model on what a human considers a “good” versus “bad” response.

Dataiku provides a complete environment for RLHF. For data collection, users can leverage AgentHub and native labeling features to create preference datasets directly from human feedback.

Typically, this data is formatted into three columns:

  • Prompt: The initial instruction or question.

  • Chosen: The response preferred by the human labeler.

  • Rejected: The less favorable response.

RLHF in Dataiku#

Once the data is collected, you can align your model using the trl (Transformer Reinforcement Learning) library and the Dataiku API. Using the Direct Preference Optimization (DPO) algorithm is often preferred over traditional Proximal Policy Optimization (PPO) as it is more stable and computationally efficient. The trl provides the DPOTrainer for that.

Note

For a complete code example of implementing DPO in Dataiku, please refer to our dedicated Fine-Tuning Code Samples section in the Developer Guide.

Once fine-tuned, use the create_finetuned_llm_version() method with your HuggingFace connection name as input. This saves the aligned model as a Fine-tuned LLM Saved Model version, inheriting from the connection’s configuration.

Reinforcement Learning from AI Feedback (RLAIF)#

While RLHF is highly effective, gathering human preference data is time-consuming and expensive. RLAIF solves this bottleneck by leveraging a “judge” LLM to generate synthetic feedback, drastically reducing the need for manual human labeling.

Orchestrating RLAIF in Dataiku#

Dataiku enables RLAIF by utilizing the LLM Mesh to orchestrate critique and revision workflows.

  1. Synthetic Generation: Use a Visual Prompt recipe to generate multiple candidate responses from your target model (here, it is the `` model from the HuggingFace connection).

  2. AI Judging: Use a more powerful, aligned model (e.g., GPT-5 or Claude 4 Sonnet) via the LLM Mesh to evaluate and rank the candidates based on specific criteria (e.g., helpfulness, lack of toxicity, etc).

  3. Alignment: The resulting dataset is then used to align the target model via DPO, exactly as you would with RLHF.

Reinforcement Learning with Verifiable Rewards (RLVR)#

For complex reasoning tasks, such as solving mathematical equations, writing functional code, or applying strict formatting rules, preference ranking is insufficient. These tasks require objective verification.

The RLVR method allows users to define strict, hard-coded reward functions to score a model, making it the right alignment method for complex & verifiable tasks. Instead of a human or an AI assessing if a response is good, a rule-based logic system (like a code compiler or a math verifier) objectively scores the output.

Implementing RLVR with GRPO in Dataiku#

To train models on complex reasoning chains, we can use the Group Relative Policy Optimization (GRPO) algorithm via the trl library and Dataiku API. This algorithm was introduced in 2024 (Shao et al., 2024), and later popularized by the DeepSeek-R1 reasoning models (DeepSeek-AI, Guo et al., 2025).

Below is an example of how to implement an RLVR workflow in a Dataiku Python recipe. The inputs are gsm8k_train and gsm8k_val, we pass both datasets to the trainer so we can track evaluation metrics during the run. The output is a Fine-tuned LLM Saved Model named rlvr_aligned_model.

We’ll break this down into three steps: preparing the model, defining the rewards, and executing the training loop. A fourth section simply explains how to visually test the alignment using a Prompt recipe.

1. Model & tokenizer preparation#

First, we load our base model. Here, we are using Qwen/Qwen2.5-1.5B-Instruct, a small model, to illustrate how it works. We are also using 4-bit quantization to fit comfortably on a single GPU.

Code 2: model_preparation.py#
import dataiku
from datasets import Dataset
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_name = "Qwen/Qwen2.5-1.5B-Instruct"
connection_name = "a_huggingface_connection_name"

# Load Datasets from the Flow
train_dataset = Dataset.from_pandas(dataiku.Dataset("gsm8k_train").get_dataframe())
val_dataset = Dataset.from_pandas(dataiku.Dataset("gsm8k_val").get_dataframe())

# The output Saved Model
saved_model = dataiku.Model("my_model_id")

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    quantization_config=quantization_config,
    use_cache=False 
)

tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

2. Defining verifiable rewards#

Next, we define our rule-based logic. We create two functions: one that checks if the extracted answer exactly matches our ground_truth, and another that checks if the model successfully used the XML tags.

Code 3: verifiable_rewards.py#
import re

def extract_xml_answer(text: str) -> str:
    """Helper to extract the answer from XML tags"""
    answer = text.split("<answer>")[-1]
    answer = answer.split("</answer>")[0]
    return answer.strip()

def correctness_reward_func(prompts, completions, ground_truth, **kwargs) -> list[float]:
    """Awards 2.0 points if the generated answer exactly matches the ground truth."""
    responses = [completion[0]['content'] for completion in completions]
    extracted_responses = [extract_xml_answer(r) for r in responses]
    return [2.0 if str(r) == str(gt) else 0.0 for r, gt in zip(extracted_responses, ground_truth)]

def format_reward_func(completions, **kwargs) -> list[float]:
    """Awards 1.0 point if the model strictly follows the <think> and <answer> XML format."""
    pattern = r"^<think>\n.*?\n</think>\n<answer>\n.*?\n</answer>\n$"
    responses = [completion[0]['content'] for completion in completions]
    return [1.0 if re.match(pattern, r, re.DOTALL) else 0.0 for r in responses]

3. Fine-tuning the model#

Finally, we pass our dataset, model, and reward functions to the GRPOTrainer. Using the Dataiku API, we ensure the newly aligned model version is seamlessly saved back to the project, as a new Saved Model version.

Code 4: model_finetuning.py#
from trl import GRPOTrainer, GRPOConfig
from peft import LoraConfig

# Create a fine-tuned LLM version using the Dataiku API
with saved_model.create_finetuned_llm_version(connection_name) as finetuned_llm_version:

    peft_config = LoraConfig(
        r=16,
        lora_alpha=32,
        lora_dropout=0.05,
        target_modules="all-linear",
        task_type="CAUSAL_LM",
    )
    
    # Define GRPO training parameters
    training_args = GRPOConfig(
        per_device_train_batch_size=4,
        num_train_epochs=1,
        evaluation_strategy="steps", # Evaluate during training
        eval_steps=50,
        output_dir=finetuned_llm_version.working_directory,
        gradient_checkpointing=True
    )

    grpo_trainer = GRPOTrainer(
        model=model,
        reward_funcs=[correctness_reward_func, format_reward_func], 
        peft_config=peft_config,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=val_dataset, # Passing validation dataset
        processing_class=tokenizer,
    )

    # Fine-tune the model using verifiable rewards
    grpo_trainer.train()
    
    # Save the model and log metadata back to Dataiku
    grpo_trainer.save_model()
    config = finetuned_llm_version.config
    config["batchSize"] = grpo_trainer.state.train_batch_size
    config["eventLog"] = grpo_trainer.state.log_history

By completing the above steps, you apply the GRPO algorithm to train models on complex reasoning chains.

4: Visual evaluation on the test set#

Once your model finishes training and is registered as a fine-tuned model in the LLM Mesh, you can easily evaluate its new reasoning capabilities! This can done in code, but using Prompt recipes works just as well.

Because we safely held out gsm8k_test from the training process, we can use it to objectively compare performance:

  1. Create a Prompt recipe using your base model (e.g., Qwen 2.5 1.5B) on the gsm8k_test dataset.

  2. Create a second Prompt recipe using your new rlvr_aligned_model on the same gsm8k_test dataset.

By comparing the outputs side-by-side, you should clearly see how the aligned model now actively “thinks” through the math problems step-by-step and strictly outputs its final answer in the correct XML format, vastly improving its accuracy.