DPO Customization

1
import json
2
import os
3
import time
4
import uuid
5
from pathlib import Path
6
from nemo_platform import NeMoPlatform, ConflictError
7
from nemo_platform.types.secrets import PlatformSecretResponse
8
from nemo_platform.types.files import HuggingfaceStorageConfigParam
9
from nemo_rl_plugin.schema import RlJobInput
10
11
12
def max_wait_time_checker(seconds: int, label: str = ""):
13
    """Return a check() that raises TimeoutError once `seconds` have elapsed."""
14
    start = time.time()
15
16
    def check():
17
        if time.time() - start > seconds:
18
            raise TimeoutError(f"{label} took longer than {seconds} seconds")
19
20
    return check
21
22
23
NMP_BASE_URL = os.environ.get("NMP_BASE_URL", "http://localhost:8080")
24
sdk = NeMoPlatform(base_url=NMP_BASE_URL, workspace="default")

1
from datasets import load_dataset, Dataset
2
3
print("Loading dataset nvidia/HelpSteer3 (preference subset)")
4
ds = load_dataset("nvidia/HelpSteer3", "preference")
5
6
# Small subsets keep the tutorial fast; larger sets train better but take longer.
7
training_size = 3000
8
validation_size = 300
9
DATASET_NAME = "dpo-dataset"
10
DATASET_PATH = Path("dpo-dataset").absolute()
11
os.makedirs(DATASET_PATH, exist_ok=True)
12
13
train_dataset = ds["train"]
14
validation_dataset = ds["validation"]
15
assert isinstance(train_dataset, Dataset) and isinstance(validation_dataset, Dataset)
16
17
# Save raw HelpSteer3 rows directly — no conversion. The platform detects the
18
# HelpSteer3 schema from the row keys (context / response1 / response2 / overall_preference).
19
train_dataset.select(range(training_size)).to_json(f"{DATASET_PATH}/training.jsonl")
20
validation_dataset.select(range(validation_size)).to_json(f"{DATASET_PATH}/validation.jsonl")
21
22
print(f"Saved training.jsonl ({training_size} rows) and validation.jsonl ({validation_size} rows)")
23
with open(f"{DATASET_PATH}/training.jsonl") as f:
24
    sample = json.loads(f.readline())
25
print("Sample keys:", sorted(sample.keys()))
26
print("overall_preference:", sample["overall_preference"])

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

1
HF_TOKEN = os.getenv("HF_TOKEN")
2
if not HF_TOKEN:
3
    raise RuntimeError("Set HF_TOKEN before running this tutorial.")
4
5
def create_or_get_secret(name: str, value: str, label: str) -> PlatformSecretResponse:
6
    try:
7
        secret = sdk.secrets.create(name=name, workspace="default", value=value)
8
        print(f"Created secret: {name}")
9
        return secret
10
    except ConflictError:
11
        print(f"Secret '{name}' already exists, continuing...")
12
        return sdk.secrets.retrieve(name=name, workspace="default")
13
14
15
hf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")

1
HF_REPO_ID = "meta-llama/Llama-3.2-1B-Instruct"
2
MODEL_NAME = "llama-3-2-1b-instruct"
3
4
storage = HuggingfaceStorageConfigParam(
5
    type="huggingface",
6
    repo_id=HF_REPO_ID,
7
    repo_type="model",
8
    token_secret=hf_secret.name,
9
)
10
11
try:
12
    base_model_fs = sdk.files.filesets.create(
13
        workspace="default", name=MODEL_NAME, description="Llama 3.2 1B Instruct base model", storage=storage
14
    )
15
    print(f"Created base model fileset: {MODEL_NAME}")
16
except ConflictError:
17
    base_model_fs = sdk.files.filesets.retrieve(workspace="default", name=MODEL_NAME)
18
    print("Base model fileset already exists.")
19
20
try:
21
    base_model = sdk.models.create(workspace="default", name=MODEL_NAME, fileset=f"default/{MODEL_NAME}")
22
except ConflictError:
23
    base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)
24
25
print(f"Base model fileset: fileset://default/{base_model.name}")
26
27
# Wait for the ModelSpec to be inferred from the checkpoint.
28
check = max_wait_time_checker(600, "Model spec")
29
while not base_model.spec:
30
    check()
31
    time.sleep(10)
32
    base_model = sdk.models.retrieve(workspace="default", name=MODEL_NAME)
33
print("Model spec ready")

1
job_suffix = uuid.uuid4().hex[:8]
2
OUTPUT_NAME = f"llama-3-2-1b-dpo-{job_suffix}"
3
4
spec = RlJobInput(
5
    model=f"default/{base_model.name}",
6
    dataset=f"default/{DATASET_NAME}",
7
    training={
8
        "type": "dpo",
9
        "epochs": 1,
10
        "batch_size": 16,
11
        "micro_batch_size": 1,
12
        "learning_rate": 5e-6,
13
        "max_seq_length": 4096,
14
        "ref_policy_kl_penalty": 0.1,
15
        "parallelism": {
16
            "num_nodes": 1,
17
            "num_gpus_per_node": 1,
18
            "tensor_parallel_size": 1,
19
            "pipeline_parallel_size": 1,
20
        },
21
    },
22
    output={"name": OUTPUT_NAME},
23
)
24
25
# `rl` auto-generates the job id (rl-<hex>); do not pass name=.
26
job = sdk.customization.rl.jobs.create(spec=spec, workspace="default")
27
print(f"Job ID: {job.job.name}")
28
print(f"Output model: {OUTPUT_NAME}")

1
from IPython.display import clear_output
2
3
check = max_wait_time_checker(7200, "DPO job")
4
while True:
5
    check()
6
    status = sdk.jobs.get_status(name=job.job.name, workspace="default")
7
    clear_output(wait=True)
8
    print(f"Job Status: {status.status}")
9
10
    step = max_steps = phase = None
11
    for job_step in status.steps or []:
12
        if job_step.name == "dpo-training":
13
            for task in job_step.tasks or []:
14
                d = task.status_details or {}
15
                step, max_steps, phase = d.get("step"), d.get("max_steps"), d.get("phase")
16
                break
17
            break
18
    if step is not None and max_steps:
19
        print(f"Training: Step {step}/{max_steps} ({100 * step / max_steps:.1f}%)")
20
        if phase:
21
            print(f"Phase: {phase}")
22
23
    if status.status in ("completed", "failed", "cancelled", "error"):
24
        print(f"\nJob finished: {status.status}")
25
        break
26
    time.sleep(15)
27
28
assert status.status == "completed"

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

1
deploy_suffix = uuid.uuid4().hex[:8]
2
DEPLOYMENT_CONFIG_NAME = f"dpo-deployment-cfg-{deploy_suffix}"
3
DEPLOYMENT_NAME = f"dpo-deployment-{deploy_suffix}"
4
5
deployment_config = sdk.inference.deployment_configs.create(
6
    workspace="default",
7
    name=DEPLOYMENT_CONFIG_NAME,
8
    engine="vllm",
9
    model_spec={"model_namespace": "default", "model_name": OUTPUT_NAME},
10
    executor_config={"gpu": 1, "image_name": "vllm/vllm-openai", "image_tag": "v0.22.1"},
11
)
12
13
deployment = sdk.inference.deployments.create(
14
    workspace="default", name=DEPLOYMENT_NAME, config=deployment_config.name
15
)
16
print(f"Deployment name: {deployment.name}")

1
check = max_wait_time_checker(1800, "Deployment")
2
while True:
3
    check()
4
    deployment_status = sdk.inference.deployments.retrieve(name=deployment.name, workspace="default")
5
    clear_output(wait=True)
6
    print(f"Deployment status: {deployment_status.status}")
7
    deployment_state = str(deployment_status.status).lower()
8
    if deployment_state in ("ready", "running"):
9
        if not sdk.models.wait_for_gateway(deployment.name, workspace="default", timeout=60):
10
            raise RuntimeError("Inference gateway did not become ready")
11
        break
12
    if deployment_state in ("failed", "error", "terminated", "lost"):
13
        raise RuntimeError(f"Deployment failed with status: {deployment_status.status}")
14
    time.sleep(15)

1
messages = [
2
    {"role": "system", "content": "You are a helpful assistant."},
3
    {"role": "user", "content": "Write a short, friendly email to a colleague asking to reschedule our meeting to Thursday."},
4
]
5
6
response = sdk.inference.gateway.provider.post(
7
    "v1/chat/completions",
8
    name=deployment.name,
9
    workspace="default",
10
    body={"model": f"default/{OUTPUT_NAME}", "messages": messages, "temperature": 0.7, "max_tokens": 256},
11
)
12
print("Model output:\n")
13
print(response["choices"][0]["message"]["content"])