diff --git a/roll/pipeline/dpo/dpo_pipeline.py b/roll/pipeline/dpo/dpo_pipeline.py index 6aeebfaf6..b6d90b4e2 100644 --- a/roll/pipeline/dpo/dpo_pipeline.py +++ b/roll/pipeline/dpo/dpo_pipeline.py @@ -1,14 +1,12 @@ import copy import json import os -import time from typing import Any, Dict, List import datasets import ray import torch from codetiming import Timer -from ray.util.timer import _Timer from torch.utils.data import DataLoader from tqdm import tqdm @@ -107,7 +105,9 @@ def __init__(self, pipeline_config: DPOConfig): if self.pipeline_config.global_template else self.pipeline_config.actor_train.data_args.template ) - encode_function = get_encode_function(template_name, self.tokenizer, self.pipeline_config.chosen_key, self.pipeline_config.rejected_key) + encode_function = get_encode_function( + template_name, self.tokenizer, self.pipeline_config.chosen_key, self.pipeline_config.rejected_key + ) self.dataset = preprocess_dataset( self.dataset, self.pipeline_config.sequence_length, @@ -168,10 +168,12 @@ def __init__(self, pipeline_config: DPOConfig): ) # Assert reference inference capacity is sufficient - reference_infer_global_batch_size = (self.pipeline_config.reference.infer_batch_size//2) * self.reference.dp_size - assert reference_infer_global_batch_size <= self.global_train_batch_size, ( - f"reference_infer_global_batch_size ({reference_infer_global_batch_size}) must be <= global train batch size ({self.global_train_batch_size})" - ) + reference_infer_global_batch_size = ( + self.pipeline_config.reference.infer_batch_size // 2 + ) * self.reference.dp_size + assert ( + reference_infer_global_batch_size <= self.global_train_batch_size + ), f"reference_infer_global_batch_size ({reference_infer_global_batch_size}) must be <= global train batch size ({self.global_train_batch_size})" if self.val_dataset: val_pipeline_config = copy.deepcopy(self.pipeline_config) @@ -189,9 +191,9 @@ def __init__(self, pipeline_config: DPOConfig): collate_fn=data_collator, ) - assert reference_infer_global_batch_size <= self.global_val_batch_size, ( - f"reference_infer_global_batch_size ({reference_infer_global_batch_size}) must be <= global val batch size ({self.global_val_batch_size})" - ) + assert ( + reference_infer_global_batch_size <= self.global_val_batch_size + ), f"reference_infer_global_batch_size ({reference_infer_global_batch_size}) must be <= global val batch size ({self.global_val_batch_size})" self.set_checkpoint_clusters(self.actor_train) @@ -219,8 +221,15 @@ def run(self): with Timer(name="step_total", logger=None) as step_total_timer: batch_dict: Dict batch: DataProto = DataProto.from_single_dict(batch_dict) - batch.meta_info = {"global_step": global_step, "is_offload_states": self.pipeline_config.is_offload_states, - "is_offload_optimizer_states_in_train_step": self.pipeline_config.is_offload_optimizer_states_in_train_step, 'loss_mask_keys': []} + batch.meta_info = { + "global_step": global_step, + "is_offload_states": self.pipeline_config.is_offload_states, + "is_offload_optimizer_states_in_train_step": self.pipeline_config.is_offload_optimizer_states_in_train_step, + "loss_mask_keys": [], + "ipo": self.pipeline_config.ipo, + "beta": self.pipeline_config.beta, + "label_smoothing": self.pipeline_config.label_smoothing, + } with Timer(name="cal_ref_log_probs", logger=None) as cal_ref_log_probs_timer: ref_log_probs = self.reference.compute_log_probs(batch, blocking=True) @@ -257,7 +266,7 @@ def run(self): global_step += 1 if global_step >= self.pipeline_config.max_steps: break - + if global_step >= self.pipeline_config.max_steps: break @@ -269,8 +278,13 @@ def val(self): for batch_dict in tqdm(self.val_dataloader): batch_dict: Dict batch: DataProto = DataProto.from_single_dict(batch_dict) - batch.meta_info = {"is_offload_states": self.pipeline_config.is_offload_states, - 'loss_mask_keys': []} + batch.meta_info = { + "is_offload_states": self.pipeline_config.is_offload_states, + "loss_mask_keys": [], + "ipo": self.pipeline_config.ipo, + "beta": self.pipeline_config.beta, + "label_smoothing": self.pipeline_config.label_smoothing, + } with Timer(name="cal_ref_log_probs", logger=None) as cal_ref_log_probs_timer: ref_log_probs = self.reference.compute_log_probs(batch, blocking=True) diff --git a/tests/pipeline/dpo/test_objective_config.py b/tests/pipeline/dpo/test_objective_config.py new file mode 100644 index 000000000..20e8254c4 --- /dev/null +++ b/tests/pipeline/dpo/test_objective_config.py @@ -0,0 +1,108 @@ +"""Exercise objective configuration through the DPO driver and worker loss on CPU.""" + +import math +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from roll.distributed.scheduler.protocol import DataProto +from roll.pipeline.dpo.actor_worker import ActorWorker +from roll.pipeline.dpo.dpo_pipeline import DPOPipeline + + +class _LocalActor: + """Replace remote/model execution while retaining the real DPO worker loss.""" + + def __init__(self): + self.theta = torch.tensor(1.0, requires_grad=True) + self.token_logps = torch.tensor([[-0.5, -0.5], [-1.0, -1.0]]) + self.worker = SimpleNamespace(strategy=SimpleNamespace(op_compute_log_probs=lambda **kw: kw["logits"])) + self.loss = None + self.gradient = None + + def compute_log_probs(self, data, blocking): + return DataProto.from_dict(tensors={"log_probs": self.theta.detach() * self.token_logps}) + + def train_step(self, data, blocking): + # The driver runs under no_grad; real training happens in another worker. + with torch.enable_grad(): + loss, metrics = ActorWorker.loss_func(self.worker, data, self.theta * self.token_logps) + loss.backward() + self.loss = loss.item() + self.gradient = self.theta.grad.item() + return DataProto(meta_info={"metrics": metrics}) + + +@pytest.fixture +def pipeline(monkeypatch): + pipeline = DPOPipeline.__new__(DPOPipeline) + pipeline.actor_train = _LocalActor() + pipeline.reference = SimpleNamespace( + compute_log_probs=lambda data, blocking: DataProto.from_dict(tensors={"log_probs": torch.zeros(2, 2)}) + ) + batch = { + "input_ids": torch.zeros((2, 3), dtype=torch.long), + "attention_mask": torch.ones((2, 3), dtype=torch.long), + "prompt_id_lens": torch.ones(2, dtype=torch.long), + } + pipeline.dataloader = [batch] + pipeline.val_dataloader = [batch] + pipeline.val_dataset = True + pipeline.state = SimpleNamespace(step=-1, log_history=[]) + pipeline.tracker = Mock() + pipeline.do_checkpoint = Mock() + # Local cluster methods return materialized DataProto objects, not Ray refs. + monkeypatch.setattr(DataProto, "materialize_concat", staticmethod(lambda data_refs: data_refs)) + monkeypatch.delenv("RAY_PROFILING", raising=False) + return pipeline + + +@pytest.mark.parametrize( + "ipo,beta,label_smoothing", + [ + pytest.param(False, 0.1, 0.0, id="default-dpo"), + pytest.param(False, 0.7, 0.0, id="nondefault-beta"), + pytest.param(False, 0.1, 0.2, id="label-smoothing"), + pytest.param(True, 0.1, 0.0, id="ipo"), + pytest.param(True, 0.7, 0.0, id="ipo-nondefault-beta"), + ], +) +def test_objective_config_controls_training_and_validation(pipeline, ipo, beta, label_smoothing): + pipeline.pipeline_config = SimpleNamespace( + ipo=ipo, + beta=beta, + label_smoothing=label_smoothing, + actor_train=SimpleNamespace(training_args=SimpleNamespace(num_train_epochs=1)), + is_offload_states=False, + is_offload_optimizer_states_in_train_step=False, + eval_steps=1, + logging_steps=1, + max_steps=1, + ) + + # One pair has policy log-ratio theta=1 and reference log-ratio zero. + # Compute the oracle independently of the production loss function. + if ipo: + expected_loss = (1 - 1 / (2 * beta)) ** 2 + expected_gradient = 2 * (1 - 1 / (2 * beta)) + else: + expected_loss = (1 - label_smoothing) * math.log1p(math.exp(-beta)) + label_smoothing * math.log1p( + math.exp(beta) + ) + expected_gradient = beta * (1 / (1 + math.exp(-beta)) - 1 + label_smoothing) + + # run() also calls the real val() before the training step. + pipeline.run() + + assert pipeline.actor_train.loss == pytest.approx(expected_loss) + assert pipeline.actor_train.gradient == pytest.approx(expected_gradient) + pipeline.tracker.log.assert_called_once() + metrics = pipeline.tracker.log.call_args.kwargs["values"] + assert metrics["actor/loss"] == pytest.approx(expected_loss) + assert metrics["val/actor/loss"] == pytest.approx(expected_loss) + for prefix in ("actor", "val/actor"): + assert metrics[f"{prefix}/chosen_reward"] == pytest.approx(-beta) + assert metrics[f"{prefix}/reject_reward"] == pytest.approx(-2 * beta) + assert metrics[f"{prefix}/acc"] == 1.0