DocumentationFine-tune ModelsTutorials

DPO Customization

View as Markdown

Run in Google Colab

Learn how to use the NeMo Platform to align a model with DPO (Direct Preference Optimization) on a preference dataset. For each prompt, DPO trains on a chosen (preferred) and a rejected response so the model prefers the chosen style — no separate reward model required.

This tutorial uses the rl customization backend (powered by NVIDIA NeMo-RL), which runs DPO on a Ray cluster. Unlike the SFT and LoRA tutorials (Docker GPU jobs), rl requires a Kubernetes-backed NeMo Platform. DPO here is full-weight (no LoRA/adapter); the output is a full model entity.

Time to complete: approximately 45-60 minutes. Job duration increases with model and dataset size.

Prerequisites

Before starting this tutorial, ensure you have:

  1. Completed the Quickstart to install the NeMo Platform and Python SDK.
  2. Installed the Python SDK (PyPI wrapper: pip install "nemo-platform[all,nemo-rl-plugin]" so RlJobInput is available; source checkout: run make bootstrap from the repository root). nemo-platform[all] does not include the RL plugin.
  3. Installed the datasets package: pip install datasets.
  4. A platform configured with platform.runtime: kubernetes. The rl (DPO) backend provisions a Ray cluster and has no local Docker fallback — submit fails fast on a Docker-runtime platform. Multi-node jobs (parallelism.num_nodes > 1) additionally require the platform-side NMP_RL_MULTINODE_SHARED_STORAGE_PATH.
  5. A Hugging Face token with access to the gated base model (this tutorial uses meta-llama/Llama-3.2-1B-Instruct). Export it as HF_TOKEN.
  6. At least one GPU with CUDA 13+ and a GPU execution profile (nemo jobs list-execution-profiles).

Quick Start

1. Initialize the SDK

The SDK needs your NeMo Platform server URL. By default http://localhost:8080 is used; set NMP_BASE_URL to override:

export NMP_BASE_URL=<YOUR_NMP_BASE_URL>
import json
import os
import time
import uuid
from pathlib import Path
from nemo_platform import NeMoPlatform, ConflictError
from nemo_platform.types.secrets import PlatformSecretResponse
from nemo_platform.types.files import HuggingfaceStorageConfigParam
from nemo_rl_plugin.schema import RlJobInput
def max_wait_time_checker(seconds: int, label: str = ""):
    """Return a check() that raises TimeoutError once `seconds` have elapsed."""
    start = time.time()
    def check():
        if time.time() - start > seconds:
            raise TimeoutError(f"{label} took longer than {seconds} seconds")
    return check
NMP_BASE_URL = os.environ.get("NMP_BASE_URL", "http://localhost:8080")
sdk = NeMoPlatform(base_url=NMP_BASE_URL, workspace="default")

2. Prepare the Preference Dataset

DPO trains on preference data. The rl backend takes a single dataset fileset that holds both training.jsonl and validation.jsonl, and auto-detects the row schema from the first line. Three preference formats are supported (see the platform’s BinaryPreferenceDatasetItemSchema / HelpSteer3DatasetItemSchema / Tulu3PreferenceDatasetItemSchema):

Binary Preference Format

Simple prompt / chosen / rejected (the prompt may be a string or a list of chat messages):

{"prompt": "What is the capital of France?", "chosen": "The capital of France is Paris.", "rejected": "I'm not sure."}

HelpSteer3 Format (used here)

A conversation context (string or chat messages), two candidate response1 / response2, and a signed overall_preference in -3..3 — negative means response 1 is preferred, positive means response 2, 0 is a tie. This is the raw schema of nvidia/HelpSteer3, so no conversion is needed:

{"context": [{"role": "user", "content": "Explain how to use git rebase"}], "response1": "...", "response2": "...", "overall_preference": -2}

Tulu3 Preference Format

Full chat conversations for both the chosen and rejected branches (each a list of messages ending with the assistant turn):

{"chosen": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "preferred"}], "rejected": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "dispreferred"}]}

Download nvidia/HelpSteer3

We use nvidia/HelpSteer3 (the preference subset), NVIDIA’s open preference dataset. It ships native train and validation splits and matches the HelpSteer3 schema above, so we upload the rows as-is — the platform’s HelpSteer3Dataset loader handles the overall_preference semantics (including ties) at training time.

from datasets import load_dataset, Dataset
print("Loading dataset nvidia/HelpSteer3 (preference subset)")
ds = load_dataset("nvidia/HelpSteer3", "preference")
# Small subsets keep the tutorial fast; larger sets train better but take longer.
training_size = 3000
validation_size = 300
DATASET_NAME = "dpo-dataset"
DATASET_PATH = Path("dpo-dataset").absolute()
os.makedirs(DATASET_PATH, exist_ok=True)
train_dataset = ds["train"]
validation_dataset = ds["validation"]
assert isinstance(train_dataset, Dataset) and isinstance(validation_dataset, Dataset)
# Save raw HelpSteer3 rows directly — no conversion. The platform detects the
# HelpSteer3 schema from the row keys (context / response1 / response2 / overall_preference).
train_dataset.select(range(training_size)).to_json(f"{DATASET_PATH}/training.jsonl")
validation_dataset.select(range(validation_size)).to_json(f"{DATASET_PATH}/validation.jsonl")
print(f"Saved training.jsonl ({training_size} rows) and validation.jsonl ({validation_size} rows)")
with open(f"{DATASET_PATH}/training.jsonl") as f:
    sample = json.loads(f.readline())
print("Sample keys:", sorted(sample.keys()))
print("overall_preference:", sample["overall_preference"])

3. Create FileSet and Upload Preference Data

Upload both JSONL files to a single FileSet so the DPO job can read them.

try:
    sdk.files.filesets.create(workspace="default", name=DATASET_NAME, description="DPO preference data")
    print(f"Created fileset: {DATASET_NAME}")
except ConflictError:
    print(f"Fileset '{DATASET_NAME}' already exists, continuing...")
sdk.files.upload(local_path=DATASET_PATH, remote_path="", fileset=DATASET_NAME, workspace="default")
print("Preference data:")
print(json.dumps([f.model_dump() for f in sdk.files.list(fileset=DATASET_NAME, workspace="default").data], indent=2, default=str))

4. Secrets Setup

The base model (meta-llama/Llama-3.2-1B-Instruct) is gated, so store your Hugging Face token as a platform secret named hf-token and reference it on the model fileset.

HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise RuntimeError("Set HF_TOKEN before running this tutorial.")
def create_or_get_secret(name: str, value: str, label: str) -> PlatformSecretResponse:
    try:
        secret = sdk.secrets.create(name=name, workspace="default", value=value)
        print(f"Created secret: {name}")
        return secret
    except ConflictError:
        print(f"Secret '{name}' already exists, continuing...")
        return sdk.secrets.retrieve(name=name, workspace="default")
hf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")

5. Create Base Model FileSet and Model Entity

DPO starts from an instruction-tuned base model. The model entity’s spec is inferred asynchronously after creation.

HF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"
MODEL_NAME = "llama-3-2-1b-instruct"
storage = HuggingfaceStorageConfigParam(
    type="huggingface",
    repo_id=HF_REPO_ID,
    repo_type="model",
    token_secret=hf_secret.name,
)
try:
    base_model_fs = sdk.files.filesets.create(
        workspace="default", name=MODEL_NAME, description="Llama 3.2 1B Instruct base model", storage=storage
    )
    print(f"Created base model fileset: {MODEL_NAME}")
except ConflictError:
    base_model_fs = sdk.files.filesets.retrieve(workspace="default", name=MODEL_NAME)
    print("Base model fileset already exists.")
try:
    base_model = sdk.models.create(workspace="default", name=MODEL_NAME, fileset=f"default/{MODEL_NAME}")
except ConflictError:
    base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)
print(f"Base model fileset: fileset://default/{base_model.name}")
# Wait for the ModelSpec to be inferred from the checkpoint.
check = max_wait_time_checker(600, "Model spec")
while not base_model.spec:
    check()
    time.sleep(10)
    base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)
print("Model spec ready")

6. Create the DPO Customization Job

Submit a DPO job to the rl backend with RlJobInput. Note the DPO-specific shape:

  • model is a string ref to the model entity; dataset is a single string ref to the preference fileset (holding both files).
  • The training method is {"type": "dpo", ...} — full-weight, no finetuning_type/LoRA.
  • ref_policy_kl_penalty is β (DPO paper): how strongly the policy stays tied to the reference model.
  • rl auto-generates the job id (rl-<hex>); read it back from the response.

Other configurable knobs: optimizer_type, adam_eps, activation_checkpointing, keep_top_k, val_at_end, preference_loss_weight, sft_loss_weight. Run nemo customization rl explain for the live schema.

job_suffix = uuid.uuid4().hex[:8]
OUTPUT_NAME = f"llama-3-2-1b-dpo-{job_suffix}"
spec = RlJobInput(
    model=f"default/{base_model.name}",
    dataset=f"default/{DATASET_NAME}",
    training={
        "type": "dpo",
        "epochs": 1,
        "batch_size": 16,
        "micro_batch_size": 1,
        "learning_rate": 5e-6,
        "max_seq_length": 4096,
        "ref_policy_kl_penalty": 0.1,
        "parallelism": {
            "num_nodes": 1,
            "num_gpus_per_node": 1,
            "tensor_parallel_size": 1,
            "pipeline_parallel_size": 1,
        },
    },
    output={"name": OUTPUT_NAME},
)
# `rl` auto-generates the job id (rl-<hex>); do not pass name=.
job = sdk.customization.rl.jobs.create(spec=spec, workspace="default")
print(f"Job ID: {job.job.name}")
print(f"Output model: {OUTPUT_NAME}")

7. Track Training Progress

The DPO job runs four steps: download -> dpo-training (Ray) -> upload -> model-entity. We poll the top-level job status and surface the training step’s progress.

from IPython.display import clear_output
check = max_wait_time_checker(7200, "DPO job")
while True:
    check()
    status = sdk.jobs.get_status(name=job.job.name, workspace="default")
    clear_output(wait=True)
    print(f"Job Status: {status.status}")
    step = max_steps = phase = None
    for job_step in status.steps or []:
        if job_step.name == "dpo-training":
            for task in job_step.tasks or []:
                d = task.status_details or {}
                step, max_steps, phase = d.get("step"), d.get("max_steps"), d.get("phase")
                break
            break
    if step is not None and max_steps:
        print(f"Training: Step {step}/{max_steps} ({100 * step / max_steps:.1f}%)")
        if phase:
            print(f"Phase: {phase}")
    if status.status in ("completed", "failed", "cancelled", "error"):
        print(f"\nJob finished: {status.status}")
        break
    time.sleep(15)
assert status.status == "completed"

Interpreting DPO training metrics (in status_details.metrics):

  • loss — the DPO loss; should trend down as the policy learns to separate chosen from rejected.
  • Reward margin (chosen minus rejected reward) — should trend up: the model increasingly prefers chosen responses.
  • Validation loss — watch for divergence from training loss (overfitting). Raise ref_policy_kl_penalty (β) or add sft_loss_weight if the policy drifts too far from the reference.

8. Validate the Output Model

DPO produces a full-weight model entity (not an adapter). Confirm it was registered.

model_entity = sdk.models.retrieve(workspace="default", name=OUTPUT_NAME)
print(model_entity.model_dump_json(indent=2))

9. Deploy and Evaluate (optional)

The DPO output is a full model, so it deploys like any full-weight checkpoint (see the Full SFT tutorial for details). We deploy with vLLM and send a chat completion.

deploy_suffix = uuid.uuid4().hex[:8]
DEPLOYMENT_CONFIG_NAME = f"dpo-deployment-cfg-{deploy_suffix}"
DEPLOYMENT_NAME = f"dpo-deployment-{deploy_suffix}"
deployment_config = sdk.inference.deployment_configs.create(
    workspace="default",
    name=DEPLOYMENT_CONFIG_NAME,
    engine="vllm",
    model_spec={"model_namespace": "default", "model_name": OUTPUT_NAME},
    executor_config={"gpu": 1, "image_name": "vllm/vllm-openai", "image_tag": "v0.22.1"},
)
deployment = sdk.inference.deployments.create(
    workspace="default", name=DEPLOYMENT_NAME, config=deployment_config.name
)
print(f"Deployment name: {deployment.name}")
check = max_wait_time_checker(1800, "Deployment")
while True:
    check()
    deployment_status = sdk.inference.deployments.retrieve(name=deployment.name, workspace="default")
    clear_output(wait=True)
    print(f"Deployment status: {deployment_status.status}")
    deployment_state = str(deployment_status.status).lower()
    if deployment_state in ("ready", "running"):
        if not sdk.models.wait_for_gateway(deployment.name, workspace="default", timeout=60):
            raise RuntimeError("Inference gateway did not become ready")
        break
    if deployment_state in ("failed", "error", "terminated", "lost"):
        raise RuntimeError(f"Deployment failed with status: {deployment_status.status}")
    time.sleep(15)
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Write a short, friendly email to a colleague asking to reschedule our meeting to Thursday."},
]
response = sdk.inference.gateway.provider.post(
    "v1/chat/completions",
    name=deployment.name,
    workspace="default",
    body={"model": f"default/{OUTPUT_NAME}", "messages": messages, "temperature": 0.7, "max_tokens": 256},
)
print("Model output:\n")
print(response["choices"][0]["message"]["content"])

Conclusion

You aligned a base model with DPO on the NeMo Platform using the rl backend:

  • Uploaded a HelpSteer3 preference dataset as-is (the platform detects the schema natively).
  • Submitted a full-weight DPO job that ran on a Ray cluster via the Kubernetes executor.
  • Registered the output as a full model entity and (optionally) deployed it for inference.

Next steps: tune the alignment strength with ref_policy_kl_penalty (β), add sft_loss_weight to anchor the policy to the chosen responses, enable activation_checkpointing for memory headroom, or scale up with parallelism. See the Training Configuration reference for the full hyperparameter set.