From 88c8af4700391e971622dc129b6f1d8a46160f24 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 18:30:08 +0000
Subject: [PATCH 01/16] Initial plan


From 083507719d1d0f6627de728d10d0e800976b8a94 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 18:37:07 +0000
Subject: [PATCH 02/16] Fix continual DPO trainer lifecycle

---
 benchmarks/dpo/README.md                      | 198 +++++++---
 .../accelerate_configs/deepspeed_zero3.yaml   |   4 +-
 .../deepspeed_zero3_offload.yaml              |  22 ++
 benchmarks/dpo/continual_dpo_trainer.py       | 355 ++++++++++--------
 benchmarks/dpo/dpo_continual.py               | 227 +++++++----
 benchmarks/dpo/sweep_configs/dpo_sweep.yaml   |  18 +-
 .../test_dpo_benchmark_regressions.py         | 121 ++++++
 7 files changed, 672 insertions(+), 273 deletions(-)
 create mode 100644 benchmarks/dpo/accelerate_configs/deepspeed_zero3_offload.yaml
 create mode 100644 test/test_benchmarks/test_dpo_benchmark_regressions.py

diff --git a/benchmarks/dpo/README.md b/benchmarks/dpo/README.md
index 53cf7a48..08643f07 100644
--- a/benchmarks/dpo/README.md
+++ b/benchmarks/dpo/README.md
@@ -1,92 +1,206 @@
-# Adaptation of TRL for Continual Learning
+# Continual DPO benchmark
 
-### Sync additional dependencies
+## Install benchmark dependencies
 
 ```sh
-uv sync --group benchmarks
+pip install -e .
+pip install pytest ruff accelerate==0.34.2 deepspeed==0.16.3 datasets>=3.2.0 numpy pandas wandb transformers>=4.49.0 trl>=0.15.2 peft>=0.14.0
 ```
 
-## Run DPO
+If you use the project's preferred workflow, sync the benchmark group instead:
 
-### Lora
+```sh
+uv sync --group benchmarks --group dev
+```
+
+## What changed in this benchmark
+
+- The continual DPO pipeline now keeps **one trainer/model lifecycle** and swaps task datasets safely instead of reusing a shared `Accelerator` across multiple trainers.
+- Reward-model policy evaluation and sampled completion logging are now **explicit task-end operations** via `--eval_policy_metrics` and `--log_completions`.
+- Reward models are **not kept on the training path** anymore; they are loaded only for explicit evaluation.
+- The default ZeRO-3 config is now a **multi-GPU template**. For a single GPU, either run without DeepSpeed first or use the offload config as a slower fallback.
+
+## Recommended sequence-length knobs
+
+DPO memory still scales with prompt/completion length and dynamic padding. For safer runs, set these explicitly:
+
+- `--max_prompt_length 256`
+- `--max_completion_length 256`
+- `--max_length 512`
+
+Increase them only after you have a stable baseline.
+
+## Fast baseline (single GPU, no DeepSpeed)
+
+Use this first to measure plain DPO throughput before adding ZeRO/offload complexity.
 
 ```sh
-uv run benchmarks/dpo/dpo_continual.py \
+python benchmarks/dpo/dpo_continual.py \
     --dataset_name benchmarks/continual_data_debug.json \
-    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
-    --reward_model_path Shahradmz/Qwen2-0.5B-Instruct_continual_data_debug_REWARD \
+    --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
+    --learning_rate 5.0e-6 \
+    --num_train_epochs 1 \
+    --per_device_train_batch_size 1 \
+    --gradient_accumulation_steps 16 \
+    --gradient_checkpointing \
+    --logging_steps 100 \
+    --eval_strategy no \
+    --bf16 \
+    --max_prompt_length 256 \
+    --max_completion_length 256 \
+    --max_length 512 \
+    --output_dir "$SCRATCH"/Qwen2.5-0.5B-DPO-baseline \
+    --no_remove_unused_columns \
+    --use_peft \
+    --lora_r 16 \
+    --lora_alpha 32
+```
+
+## Recommended memory-efficient Qwen 3B run (single GPU)
+
+For constrained hardware, prefer **QLoRA** over full fine-tuning.
+
+```sh
+python benchmarks/dpo/dpo_continual.py \
+    --dataset_name benchmarks/continual_data_debug.json \
+    --model_name_or_path Qwen/Qwen2.5-3B-Instruct \
     --learning_rate 5.0e-6 \
     --num_train_epochs 1 \
-    --per_device_train_batch_size 2 \
-    --gradient_accumulation_steps 8 \
+    --per_device_train_batch_size 1 \
+    --gradient_accumulation_steps 16 \
     --gradient_checkpointing \
-    --logging_steps 5 \
+    --logging_steps 100 \
     --eval_strategy steps \
-    --eval_steps 5 \
-    --save_steps 5 \
+    --eval_steps 200 \
     --bf16 \
-    --output_dir "$SCRATCH"/Qwen2-0.5B-DPO-test \
+    --max_prompt_length 256 \
+    --max_completion_length 256 \
+    --max_length 512 \
+    --output_dir "$SCRATCH"/Qwen2.5-3B-DPO-qlora \
     --no_remove_unused_columns \
     --use_peft \
-    --lora_r 32 \
-    --lora_alpha 16
+    --load_in_4bit \
+    --lora_r 16 \
+    --lora_alpha 32
 ```
 
-### Lora with accelerate
+Optional explicit task-end reward evaluation/logging:
+
+```sh
+    --reward_model_path Shahradmz/Qwen2-0.5B-Instruct_continual_data_debug_REWARD \
+    --eval_policy_metrics \
+    --log_completions \
+    --completion_logging_batches 1
+```
+
+Those flags add extra evaluation work on purpose; leave them off for raw throughput measurements.
+
+## Multi-GPU ZeRO-3 (throughput / memory sharding)
+
+`benchmarks/dpo/accelerate_configs/deepspeed_zero3.yaml` is a **template for multi-GPU runs**. Set `num_processes` to the number of GPUs you actually launch.
 
 ```sh
 accelerate launch --config_file benchmarks/dpo/accelerate_configs/deepspeed_zero3.yaml \
     benchmarks/dpo/dpo_continual.py \
     --dataset_name benchmarks/continual_data_debug.json \
-    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
-    --reward_model_path Shahradmz/Qwen2-0.5B-Instruct_continual_data_debug_REWARD \
+    --model_name_or_path Qwen/Qwen2.5-3B-Instruct \
     --learning_rate 5.0e-6 \
     --num_train_epochs 1 \
-    --per_device_train_batch_size 2 \
-    --gradient_accumulation_steps 8 \
+    --per_device_train_batch_size 1 \
+    --gradient_accumulation_steps 16 \
     --gradient_checkpointing \
-    --logging_steps 5 \
+    --logging_steps 100 \
     --eval_strategy steps \
-    --eval_steps 5 \
-    --save_steps 5 \
+    --eval_steps 200 \
     --bf16 \
-    --output_dir "$SCRATCH"/Qwen2-0.5B-DPO-test \
+    --max_prompt_length 256 \
+    --max_completion_length 256 \
+    --max_length 512 \
+    --output_dir "$SCRATCH"/Qwen2.5-3B-DPO-zero3 \
     --no_remove_unused_columns \
     --use_peft \
-    --lora_r 32 \
-    --lora_alpha 16 \
-    --wandb_project Qwen2-0.5B-DPO_lora_test
+    --lora_r 16 \
+    --lora_alpha 32
 ```
 
-### Full training
+Use this config only when you are actually launching multiple processes/GPUs. A one-process ZeRO-3 launch does **not** give multi-GPU parameter sharding.
+
+## Single-GPU DeepSpeed fallback
+
+If plain single-GPU QLoRA still does not fit, the slower fallback is CPU offload:
 
 ```sh
-uv run benchmarks/dpo/dpo_continual.py \
+accelerate launch --config_file benchmarks/dpo/accelerate_configs/deepspeed_zero3_offload.yaml \
+    benchmarks/dpo/dpo_continual.py \
     --dataset_name benchmarks/continual_data_debug.json \
-    --model_name_or_path Qwen/Qwen2-0.5B-Instruct \
-    --reward_model_path Shahradmz/Qwen2-0.5B-Instruct_continual_data_debug_REWARD \
+    --model_name_or_path Qwen/Qwen2.5-3B-Instruct \
+    --learning_rate 5.0e-6 \
+    --num_train_epochs 1 \
+    --per_device_train_batch_size 1 \
+    --gradient_accumulation_steps 16 \
+    --gradient_checkpointing \
+    --logging_steps 100 \
+    --eval_strategy no \
+    --bf16 \
+    --max_prompt_length 256 \
+    --max_completion_length 256 \
+    --max_length 512 \
+    --output_dir "$SCRATCH"/Qwen2.5-3B-DPO-zero3-offload \
+    --no_remove_unused_columns \
+    --use_peft \
+    --load_in_4bit \
+    --lora_r 16 \
+    --lora_alpha 32
+```
+
+This is a fallback for memory pressure, not a fast path.
+
+## Full fine-tuning
+
+Full fine-tuning keeps the policy model trainable and also needs a reference model for standard DPO. On 3B models this is much more memory intensive than LoRA/QLoRA.
+
+```sh
+python benchmarks/dpo/dpo_continual.py \
+    --dataset_name benchmarks/continual_data_debug.json \
+    --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
     --learning_rate 5.0e-7 \
     --num_train_epochs 1 \
-    --per_device_train_batch_size 2 \
-    --gradient_accumulation_steps 8 \
+    --per_device_train_batch_size 1 \
+    --gradient_accumulation_steps 16 \
     --gradient_checkpointing \
-    --logging_steps 25 \
+    --logging_steps 100 \
     --eval_strategy steps \
-    --eval_steps 50 \
-    --output_dir Qwen2-0.5B-DPO \
+    --eval_steps 200 \
+    --bf16 \
+    --max_prompt_length 256 \
+    --max_completion_length 256 \
+    --max_length 512 \
+    --output_dir "$SCRATCH"/Qwen2.5-0.5B-DPO-full \
     --no_remove_unused_columns
 ```
 
-### Run a sweep with wandb
+## LoRA vs QLoRA vs full fine-tuning
+
+- **LoRA**: `--use_peft`; keeps the base model in standard precision.
+- **QLoRA**: `--use_peft --load_in_4bit`; usually the best single-GPU choice for 3B models.
+- **Full fine-tuning**: omit PEFT flags; highest memory use.
+
+With PEFT, TRL can use the base model without a separate train-time reference-model copy, which is typically the lowest-memory DPO setup in this benchmark.
+
+## Important limitations
+
+- Hardware limits still apply. These configs reduce risk; they do **not** guarantee that OOM is impossible.
+- Reward-model policy evaluation can still require substantial memory because it temporarily loads an additional model for explicit evaluation.
+- If a run is unstable, reduce `--max_prompt_length`, `--max_completion_length`, `--max_length`, and/or increase `--gradient_accumulation_steps` before increasing batch size.
+
+## Run a sweep with wandb
 
 ```sh
-wandb sweep sweep_configs/dpo_sweep.yaml    # which returns the SWEEP_ID
+wandb sweep benchmarks/dpo/sweep_configs/dpo_sweep.yaml
 ```
 
-and
+Then run the returned sweep ID:
 
 ```sh
 wandb agent <SWEEP_ID>
 ```
-
-- All details per task and hyperparameters are going to be loaded in your wandb dashboard.
diff --git a/benchmarks/dpo/accelerate_configs/deepspeed_zero3.yaml b/benchmarks/dpo/accelerate_configs/deepspeed_zero3.yaml
index 7f17a48f..4b4312ea 100644
--- a/benchmarks/dpo/accelerate_configs/deepspeed_zero3.yaml
+++ b/benchmarks/dpo/accelerate_configs/deepspeed_zero3.yaml
@@ -2,6 +2,8 @@ compute_environment: LOCAL_MACHINE
 debug: false
 deepspeed_config:
   deepspeed_multinode_launcher: standard
+  offload_optimizer_device: none
+  offload_param_device: none
   zero3_init_flag: true
   zero3_save_16bit_model: true
   zero_stage: 3
@@ -11,7 +13,7 @@ machine_rank: 0
 main_training_function: main
 mixed_precision: bf16
 num_machines: 1
-num_processes: 1  # TODO change to whatever number of gpus is used
+num_processes: 4  # Set this to the number of GPUs you are actually launching.
 rdzv_backend: static
 same_network: true
 tpu_env: []
diff --git a/benchmarks/dpo/accelerate_configs/deepspeed_zero3_offload.yaml b/benchmarks/dpo/accelerate_configs/deepspeed_zero3_offload.yaml
new file mode 100644
index 00000000..c87f575c
--- /dev/null
+++ b/benchmarks/dpo/accelerate_configs/deepspeed_zero3_offload.yaml
@@ -0,0 +1,22 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  deepspeed_multinode_launcher: standard
+  offload_optimizer_device: cpu
+  offload_param_device: cpu
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 1
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
diff --git a/benchmarks/dpo/continual_dpo_trainer.py b/benchmarks/dpo/continual_dpo_trainer.py
index 8b68e7e6..f6dd8a28 100644
--- a/benchmarks/dpo/continual_dpo_trainer.py
+++ b/benchmarks/dpo/continual_dpo_trainer.py
@@ -1,7 +1,6 @@
-import functools
-import inspect
 import os
 from collections import defaultdict
+from contextlib import contextmanager
 from dataclasses import dataclass, field
 from typing import Callable, Optional, Union
 
@@ -9,11 +8,12 @@
 import pandas as pd
 import torch
 import torch.nn as nn
-from accelerate import Accelerator, PartialState
+from accelerate import PartialState
 from accelerate.utils import gather_object
 from datasets import Dataset
 from torch.utils.data import DataLoader
 from transformers import (
+    AutoModelForSequenceClassification,
     BaseImageProcessor,
     DataCollator,
     DataCollatorWithPadding,
@@ -28,12 +28,7 @@
 from trl import DPOTrainer, ScriptArguments
 from trl.models.utils import unwrap_model_for_generation
 from trl.trainer.dpo_config import DPOConfig
-from trl.trainer.utils import (
-    batch_generation,
-    disable_dropout_in_model,
-    get_reward,
-    prepare_deepspeed,
-)
+from trl.trainer.utils import batch_generation, disable_dropout_in_model, get_reward
 from typing_extensions import override
 
 import wandb as wb
@@ -86,7 +81,7 @@ class ContinualDPOConfig(DPOConfig):
     response_length: int = field(
         default=53,
         metadata={
-            'help': 'Length of the response. Borrowed from PPOCOnfig and used only for evaluation.'
+            'help': 'Length of the response. Borrowed from PPOConfig and used only for evaluation.'
         },
     )
     temperature: float = field(
@@ -99,13 +94,27 @@ class ContinualDPOConfig(DPOConfig):
         default=False,
         metadata={'help': 'Whether to use greedy policy for evaluation.'},
     )
+    eval_policy_metrics: bool = field(
+        default=False,
+        metadata={
+            'help': 'Run reward-model policy evaluation only during the explicit task-end evaluation path.'
+        },
+    )
+    log_completions: bool = field(
+        default=False,
+        metadata={
+            'help': 'Log sampled completions only during the explicit task-end evaluation path.'
+        },
+    )
+    completion_logging_batches: int = field(
+        default=1,
+        metadata={
+            'help': 'Number of evaluation batches to sample when --log_completions is enabled.'
+        },
+    )
 
 
 class ContinualDPOTrainer(DPOTrainer):
-    # Shared accelerator instance across all trainer instances
-    shared_accelerator: Optional[Accelerator] = None
-    accelerator: Accelerator  # now non-optional after creation
-
     @override
     def __init__(
         self,
@@ -156,83 +165,71 @@ def __init__(
             preprocess_logits_for_metrics,
             peft_config,
         )
-        # setting a reward model only for evaluation purposes
-        # The reward model setting code comes from TRL PPOTrainer https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L284
+
         self.reward_model = reward_model
-        if self.reward_model is not None:
+        if self.reward_model is not None and not isinstance(self.reward_model, str):
             disable_dropout_in_model(self.reward_model)
-        if self.is_deepspeed_enabled:
-            if self.reward_model is not None:
-                self.reward_model = prepare_deepspeed(
-                    self.reward_model,
-                    args.per_device_train_batch_size,
-                    args.fp16,
-                    args.bf16,
-                )
-        else:
-            # Ensure reward_model is a model instance (if given as a str then load it)
-            if isinstance(self.reward_model, str):
-                from transformers import AutoModelForSequenceClassification
-
-                self.reward_model = AutoModelForSequenceClassification.from_pretrained(
-                    self.reward_model
-                )
-            # Ensure accelerator is ready. It should be set already by DPOTrainer.
-            assert self.accelerator is not None, 'Accelerator must be initialized'
-            if self.reward_model is not None:
-                self.reward_model = self.reward_model.to(self.accelerator.device)
+            self.reward_model.eval()
 
         if eval_policy_dataset is not None:
             self.eval_policy_dataset = self.preprocess_policy_dataset(
                 eval_policy_dataset
             )
-            # using the same data_collator as in PPO trainer
-            data_collator = DataCollatorWithPadding(self.processing_class)
-            self.eval_policy_dataloader = DataLoader(
-                self.eval_policy_dataset,
-                batch_size=self.args.per_device_eval_batch_size,
-                collate_fn=data_collator,
-                drop_last=True,
-            )  # no need to shuffle eval dataset
-            # Ensure accelerator is available
-            # TODO remove the check once ruff issues are resolved
-            # fmt: off
-            assert self.accelerator is not None, 'Accelerator must be assigned before prepare()'
-            # fmt: on
-            self.eval_policy_dataloader = self.accelerator.prepare(
-                self.eval_policy_dataloader
+            self.eval_policy_dataloader = self._build_eval_policy_dataloader(
+                self.eval_policy_dataset
             )
-
         else:
             self.eval_policy_dataset = None
             self.eval_policy_dataloader = None
 
-    def create_accelerator_and_postprocess(self) -> None:
-        # Only initialize a new Accelerator if one does not exist
-        if ContinualDPOTrainer.shared_accelerator is None:
-            super().create_accelerator_and_postprocess()
-            ContinualDPOTrainer.shared_accelerator = self.accelerator
-        else:
-            # Reuse the shared accelerator
-            self.accelerator = ContinualDPOTrainer.shared_accelerator
-            self.gather_function = self.accelerator.gather_for_metrics
-            if (
-                'use_gather_object'
-                in inspect.signature(self.gather_function).parameters.keys()
-            ):
-                self.gather_function = functools.partial(
-                    self.gather_function,
-                    use_gather_object=self.args.eval_use_gather_object,
-                )
-            self.is_deepspeed_enabled = (
-                getattr(self.accelerator.state, 'deepspeed_plugin', None) is not None
+    def _build_eval_policy_dataloader(self, dataset: Dataset) -> DataLoader:
+        data_collator = DataCollatorWithPadding(self.processing_class)
+        dataloader = DataLoader(
+            dataset,
+            batch_size=self.args.per_device_eval_batch_size,
+            collate_fn=data_collator,
+            drop_last=True,
+        )
+        return self.accelerator.prepare(dataloader)
+
+    def set_task_datasets(
+        self,
+        train_dataset: Dataset,
+        eval_dataset: Optional[Dataset] = None,
+        dataset_name: str = 'task',
+    ) -> None:
+        self.train_dataset = self._prepare_dataset(
+            train_dataset,
+            self.processing_class,
+            self.args,
+            f'{dataset_name}-train',
+        )
+        self.eval_dataset = None
+        self.eval_policy_dataset = None
+        self.eval_policy_dataloader = None
+
+        if eval_dataset is not None:
+            self.eval_dataset = self._prepare_dataset(
+                eval_dataset,
+                self.processing_class,
+                self.args,
+                f'{dataset_name}-eval',
             )
-            self.is_fsdp_enabled = (
-                getattr(self.accelerator.state, 'fsdp_plugin', None) is not None
+            self.eval_policy_dataset = self.preprocess_policy_dataset(eval_dataset)
+            self.eval_policy_dataloader = self._build_eval_policy_dataloader(
+                self.eval_policy_dataset
             )
 
+    def set_reward_model(
+        self,
+        reward_model: Optional[Union[PreTrainedModel, nn.Module, str]],
+    ) -> None:
+        self.reward_model = reward_model
+        if self.reward_model is not None and not isinstance(self.reward_model, str):
+            disable_dropout_in_model(self.reward_model)
+            self.reward_model.eval()
+
     def preprocess_policy_dataset(self, dataset: Dataset) -> Dataset:
-        # The code is from TRL PPO script https://github.com/huggingface/trl/blob/main/examples/scripts/ppo/ppo.py
         dataset_text_field = 'prompt'
 
         def tokenize(element: dict) -> dict[str, list[int]]:
@@ -250,20 +247,56 @@ def prepare_dataset(ds: Dataset) -> Dataset:
                 num_proc=self.args.dataset_num_proc,
             )
 
-        # Compute only on main process for faster data processing.
         with PartialState().local_main_process_first():
             dataset = prepare_dataset(dataset)
         return dataset
 
-    def evaluate_policy(self) -> dict:
-        """Evaluate the policy using the evaluation policy dataloader.
+    @contextmanager
+    def reward_model_context(
+        self,
+        reward_model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,
+    ):
+        active_reward_model = reward_model if reward_model is not None else self.reward_model
+        if active_reward_model is None:
+            yield None
+            return
+
+        loaded_in_context = False
+        if isinstance(active_reward_model, str):
+            active_reward_model = AutoModelForSequenceClassification.from_pretrained(
+                active_reward_model,
+                num_labels=1,
+            )
+            loaded_in_context = True
+
+        disable_dropout_in_model(active_reward_model)
+        active_reward_model.eval()
+
+        original_device = next(active_reward_model.parameters()).device
+        target_device = self.accelerator.device
+        moved_to_accelerator = original_device != target_device
+        if moved_to_accelerator:
+            active_reward_model = active_reward_model.to(target_device)
+
+        try:
+            yield active_reward_model
+        finally:
+            if moved_to_accelerator:
+                active_reward_model = active_reward_model.to(original_device)
+            if loaded_in_context:
+                del active_reward_model
+            if torch.cuda.is_available() and (moved_to_accelerator or loaded_in_context):
+                torch.cuda.empty_cache()
+
+    def evaluate_policy(
+        self,
+        reward_model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,
+    ) -> dict:
+        """Evaluate the policy on the evaluation prompts with an optional reward model."""
+        if self.eval_policy_dataloader is None:
+            return {}
 
-        Returns:
-            dict: A dictionary containing evaluation metrics.
-        """
-        # The code is heavily based on the training loop of TRL PPOTrainer function https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L677
         mode = self.model.training
-        # there is no self.model? TODO
         self.model.eval()
         eval_metrics = defaultdict(list)
         processing_class = self.processing_class
@@ -274,7 +307,6 @@ def evaluate_policy(self) -> dict:
                 do_sample=False,
             )
         else:
-            # Using the same hyperpaprams as during PPO training
             generation_config = GenerationConfig(
                 max_new_tokens=self.args.response_length,
                 temperature=(self.args.temperature + 1e-7),
@@ -283,8 +315,12 @@ def evaluate_policy(self) -> dict:
                 do_sample=True,
             )
 
-        with torch.no_grad():
-            if self.eval_policy_dataloader is not None:
+        with self.reward_model_context(reward_model) as active_reward_model:
+            if active_reward_model is None:
+                self.model.train(mode)
+                return {}
+
+            with torch.inference_mode():
                 for batch in self.eval_policy_dataloader:
                     query = batch['input_ids'].to(self.accelerator.device)
                     context_length = query.shape[1]
@@ -301,12 +337,9 @@ def evaluate_policy(self) -> dict:
                             generation_config,
                         )
                         response = query_response[:, context_length:]
-                    postprocessed_response = response
-                    postprocessed_query_response = torch.cat(
-                        (query, postprocessed_response), 1
-                    )
+                    postprocessed_query_response = torch.cat((query, response), 1)
                     _, score, _ = get_reward(
-                        self.reward_model,
+                        active_reward_model,
                         postprocessed_query_response,
                         processing_class.pad_token_id,
                         context_length,
@@ -314,81 +347,99 @@ def evaluate_policy(self) -> dict:
                     eval_metrics['score'].extend(
                         self.accelerator.gather_for_metrics(score).float().cpu().numpy()
                     )
+
         self.model.train(mode)
+        if not eval_metrics:
+            return {}
         return {'eval_' + k: float(np.mean(v)) for k, v in eval_metrics.items()}
 
-    def log(
-        self, logs: dict[str, Union[float, dict]], start_time: Optional[float] = None
-    ) -> None:
-        """Log `logs` on the various objects watching training, including stored metrics."""
-        train_eval = 'train' if 'loss' in logs else 'eval'
-        print(f'Logging {train_eval} metrics...')
-        if train_eval == 'eval':
-            print('Computing policy metrics...')
-            eval_policy_metrics = self.evaluate_policy()
-            logs.update(eval_policy_metrics)
-
-        # TODO: Only generation sample completions every x steps
-        do_generate_completions = True
-        if do_generate_completions:
-            self._generate_completions()
-            torch.cuda.empty_cache()
-        return super().log(logs, start_time)
+    def generate_completions_table(
+        self,
+        reward_model: Optional[Union[PreTrainedModel, nn.Module, str]] = None,
+        max_batches: Optional[int] = None,
+    ) -> Optional[pd.DataFrame]:
+        """Generate a small sample of completions for explicit task-end inspection."""
+        if self.eval_dataset is None:
+            return None
 
-    def _generate_completions(self) -> None:
-        # Config from: https://github.com/huggingface/trl/blob/56e57662053e2d0cc6302dad404820b0c0ec6a91/trl/trainer/ppo_trainer.py#L688
+        eval_dataloader = self.get_eval_dataloader()
         generation_config = GenerationConfig(
-            max_new_tokens=53,
+            max_new_tokens=self.args.response_length,
             temperature=(0.01 + 1e-7),
             top_k=0.0,
             top_p=1.0,
             do_sample=True,
         )
         table = defaultdict(list)
-        with torch.no_grad():
-            with unwrap_model_for_generation(
-                self.model,
-                self.accelerator,
-                gather_deepspeed3_params=None,
-            ) as unwrapped_model:
-                for batch in self.eval_dataloader:
-                    query = batch['input_ids']
-                    context_length = query.shape[1]
-                    query_response, _ = batch_generation(
-                        unwrapped_model,
-                        query,
-                        query.shape[0],
-                        self.processing_class.pad_token_id,
-                        generation_config,
-                    )
-                    response = query_response[:, context_length:]
-                    postprocessed_response = response
-                    postprocessed_query_response = torch.cat(
-                        (query, postprocessed_response), 1
-                    )
-                    _, score, _ = get_reward(
-                        self.reward_model,
-                        postprocessed_query_response,
-                        self.processing_class.pad_token_id,
-                        context_length,
-                    )
+        batches_to_log = (
+            max_batches if max_batches is not None else self.args.completion_logging_batches
+        )
+        if batches_to_log <= 0:
+            return None
+
+        mode = self.model.training
+        self.model.eval()
 
-                    queries = gather_object(
-                        self.processing_class.batch_decode(
-                            query, skip_special_tokens=True
+        with self.reward_model_context(reward_model) as active_reward_model:
+            with torch.inference_mode():
+                with unwrap_model_for_generation(
+                    self.model,
+                    self.accelerator,
+                    gather_deepspeed3_params=None,
+                ) as unwrapped_model:
+                    for batch_index, batch in enumerate(eval_dataloader):
+                        if batch_index >= batches_to_log:
+                            break
+
+                        query = batch['prompt_input_ids'].to(self.accelerator.device)
+                        context_length = query.shape[1]
+                        query_response, _ = batch_generation(
+                            unwrapped_model,
+                            query,
+                            query.shape[0],
+                            self.processing_class.pad_token_id,
+                            generation_config,
                         )
-                    )
-                    responses = gather_object(
-                        self.processing_class.batch_decode(postprocessed_response)
-                    )
-                    scores = (
-                        self.accelerator.gather_for_metrics(score).float().cpu().numpy()
-                    )
-                    table['query'].extend(queries)
-                    table['model response'].extend(responses)
-                    table['score'].extend(scores)
-                    break
+                        response = query_response[:, context_length:]
+                        postprocessed_query_response = torch.cat((query, response), 1)
+
+                        if active_reward_model is not None:
+                            _, score, _ = get_reward(
+                                active_reward_model,
+                                postprocessed_query_response,
+                                self.processing_class.pad_token_id,
+                                context_length,
+                            )
+                            table['score'].extend(
+                                self.accelerator.gather_for_metrics(score)
+                                .float()
+                                .cpu()
+                                .numpy()
+                            )
+
+                        queries = gather_object(
+                            self.processing_class.batch_decode(
+                                query,
+                                skip_special_tokens=True,
+                            )
+                        )
+                        responses = gather_object(
+                            self.processing_class.batch_decode(
+                                response,
+                                skip_special_tokens=True,
+                            )
+                        )
+                        table['query'].extend(queries)
+                        table['model response'].extend(responses)
+
+        self.model.train(mode)
+        if not table:
+            return None
 
         df = pd.DataFrame(table)
         if self.accelerator.is_main_process and wb.run is not None:
             wb.log({'completions': wb.Table(dataframe=df)})
+        return df
+
+    def _generate_completions(self) -> None:
+        self.generate_completions_table(reward_model=self.reward_model)
diff --git a/benchmarks/dpo/dpo_continual.py b/benchmarks/dpo/dpo_continual.py
index 080d8d51..fe72faf1 100644
--- a/benchmarks/dpo/dpo_continual.py
+++ b/benchmarks/dpo/dpo_continual.py
@@ -1,13 +1,10 @@
 """Adaptation of the DPO TRL training script for continual learning."""
 
 import os
+import warnings
+from typing import Optional
 
 import torch
-from continual_dpo_trainer import (
-    ContinualDPOArguments,
-    ContinualDPOConfig,
-    ContinualDPOTrainer,
-)
 from datasets import Dataset
 from transformers import (
     AutoModelForCausalLM,
@@ -33,12 +30,95 @@
 
 
 # The code is based on TRL DPO script https://github.com/huggingface/trl/blob/main/trl/scripts/dpo.py
+
+def warn_for_memory_settings(
+    training_args: ContinualDPOConfig,
+    model_args: ModelConfig,
+) -> None:
+    if training_args.max_completion_length is None:
+        warnings.warn(
+            'max_completion_length is unset. Long chosen/rejected responses can still cause large padded DPO batches; '
+            'set it explicitly for tighter memory bounds.',
+            stacklevel=2,
+        )
+
+    if not training_args.gradient_checkpointing:
+        warnings.warn(
+            'gradient_checkpointing is disabled. For larger DPO runs this will increase activation memory pressure.',
+            stacklevel=2,
+        )
+
+    if not training_args.bf16 and not training_args.fp16:
+        warnings.warn(
+            'Neither bf16 nor fp16 is enabled. Full-precision DPO is usually much more memory intensive.',
+            stacklevel=2,
+        )
+
+    if not model_args.use_peft and training_args.reward_model_path is not None:
+        warnings.warn(
+            'Full-parameter DPO with an explicit reward-model evaluation path is the highest-memory configuration. '
+            'Prefer --use_peft (or QLoRA with --load_in_4bit) for constrained hardware.',
+            stacklevel=2,
+        )
+
+
+def get_task_reward_model_path(
+    reward_model_root: Optional[str],
+    task_index: int,
+) -> Optional[str]:
+    if reward_model_root is None:
+        return None
+    return f'{reward_model_root}_{task_index}'
+
+
+def validate_reward_model_paths(
+    reward_model_root: Optional[str],
+    num_tasks: int,
+) -> None:
+    if reward_model_root is None:
+        return
+
+    for task_index in range(num_tasks):
+        reward_path = get_task_reward_model_path(reward_model_root, task_index)
+        assert reward_path is not None
+        try:
+            AutoModelForSequenceClassification.from_pretrained(
+                reward_path,
+                num_labels=1,
+            )
+        except Exception:
+            if not os.path.exists(reward_path):
+                raise ValueError(f'Reward model not found at {reward_path}') from None
+
+
+def load_reward_model_for_task(
+    reward_model_root: Optional[str],
+    task_index: int,
+    torch_dtype: Optional[torch.dtype],
+    trust_remote_code: bool,
+) -> Optional[AutoModelForSequenceClassification]:
+    reward_path = get_task_reward_model_path(reward_model_root, task_index)
+    if reward_path is None:
+        return None
+
+    reward_model_kwargs = {
+        'num_labels': 1,
+        'trust_remote_code': trust_remote_code,
+    }
+    if isinstance(torch_dtype, torch.dtype):
+        reward_model_kwargs['torch_dtype'] = torch_dtype
+
+    return AutoModelForSequenceClassification.from_pretrained(
+        reward_path,
+        **reward_model_kwargs,
+    )
+
+
 def main(
     script_args: ContinualDPOArguments,
     training_args: ContinualDPOConfig,
     model_args: ModelConfig,
 ) -> None:
-    # Determine torch dtype and quantization configs
     torch_dtype = (
         model_args.torch_dtype
         if model_args.torch_dtype in ['auto', None]
@@ -47,9 +127,9 @@ def main(
     if script_args.wandb_run_name is not None:
         training_args.run_name = script_args.wandb_run_name
 
-    quantization_config = get_quantization_config(model_args)
+    warn_for_memory_settings(training_args, model_args)
 
-    # Model & Tokenizer Setup
+    quantization_config = get_quantization_config(model_args)
     model_kwargs = dict(
         revision=model_args.model_revision,
         attn_implementation=model_args.attn_implementation,
@@ -73,22 +153,20 @@ def main(
     else:
         ref_model = None
 
-    # Load tokenizer and set chat template if needed
     tokenizer = AutoTokenizer.from_pretrained(
-        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code
+        model_args.model_name_or_path,
+        trust_remote_code=model_args.trust_remote_code,
     )
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
     if tokenizer.chat_template is None:
         tokenizer.chat_template = SIMPLE_CHAT_TEMPLATE
 
-    # Distributed training hack
     if script_args.ignore_bias_buffers:
         model._ddp_params_and_buffers_to_ignore = [
             name for name, buffer in model.named_buffers() if buffer.dtype == torch.bool
         ]
 
-    # Initialize continual dataset
     continual_dataset: list[dict[str, Dataset]] = init_continual_dataset(
         script_args.dataset_name,
         mock=training_args.mock,
@@ -96,81 +174,84 @@ def main(
         tools=training_args.tools,
     )
     output_dir = training_args.output_dir
+    eval_enabled = training_args.eval_strategy != 'no'
+    explicit_policy_eval = training_args.eval_policy_metrics or training_args.log_completions
 
-    # check if the reward models are present either in the path or in the hub
-    if training_args.reward_model_path is not None:
-        for i in range(len(continual_dataset)):
-            reward_path = training_args.reward_model_path + '_' + str(i)
-            # first check the hub if the model is present
-            try:
-                AutoModelForSequenceClassification.from_pretrained(
-                    reward_path, num_labels=1
-                )
-            except:
-                # if not found in the hub, check the local path
-                if not os.path.exists(reward_path):
-                    raise ValueError(f'Reward model not found at {reward_path}')
-
-    # Task Loop
-    for i, dataset in enumerate(continual_dataset):
-        current_dataset_name: str = f'dataset-{i}'
+    if training_args.eval_policy_metrics and training_args.reward_model_path is None:
+        raise ValueError('--eval_policy_metrics requires --reward_model_path.')
+
+    if explicit_policy_eval:
+        validate_reward_model_paths(
+            training_args.reward_model_path,
+            len(continual_dataset),
+        )
+
+    first_dataset = continual_dataset[0]
+    trainer = ContinualDPOTrainer(
+        args=training_args,
+        processing_class=tokenizer,
+        model=model,
+        ref_model=ref_model,
+        train_dataset=first_dataset[script_args.dataset_train_split],
+        eval_dataset=first_dataset.get(script_args.dataset_test_split),
+        peft_config=peft_config,
+    )
+
+    if wb.run is not None:
+        wb.log({'dataset/name': script_args.dataset_name})
+
+    for task_index, dataset in enumerate(continual_dataset):
+        current_dataset_name = f'dataset-{task_index}'
         training_args.output_dir = f'{output_dir}/{current_dataset_name}'
+        trainer.args.output_dir = training_args.output_dir
 
-        # Load reward model if path provided
-        if training_args.reward_model_path is not None:
-            reward_model = AutoModelForSequenceClassification.from_pretrained(
-                training_args.reward_model_path + f'_{str(i)}', num_labels=1
+        if task_index > 0:
+            trainer.set_task_datasets(
+                train_dataset=dataset[script_args.dataset_train_split],
+                eval_dataset=dataset.get(script_args.dataset_test_split),
+                dataset_name=current_dataset_name,
             )
 
-        trainer = ContinualDPOTrainer(
-            args=training_args,
-            processing_class=tokenizer,
-            model=model,
-            ref_model=ref_model,
-            reward_model=reward_model
-            if training_args.reward_model_path is not None
-            else None,
-            train_dataset=dataset[script_args.dataset_train_split],
-            eval_dataset=dataset[script_args.dataset_test_split]
-            if training_args.eval_strategy != 'no'
-            else None,
-            peft_config=peft_config,
-        )
-
-        # TODO will throw Invalidate trace cache @ step 10: expected module 11, but got module 19
-        # https://github.com/deepspeedai/DeepSpeed/issues/6870
-        # Fix with deepspeed fix release
         print('Training dataset:', current_dataset_name)
         trainer.train()
 
-        if training_args.eval_strategy != 'no':
+        should_run_task_eval = eval_enabled or explicit_policy_eval
+        if should_run_task_eval and trainer.eval_dataset is not None:
             metrics = trainer.evaluate()
-            if i == 0:
-                trainer.log({'dataset': {'name': script_args.dataset_name}})
-            metrics['dataset'] = i
-            # Log evaluation metrics under a hierarchy using slashes for wandb
-            print(f'eval/dataset/{i}')
-            trainer.log_metrics(f'eval/dataset/{i}', metrics)
-            trainer.save_metrics(f'eval', metrics)
-            wb.log({'eval': {'last': metrics}})  # type: ignore[attr-defined]
-            wb.log({f'task/{current_dataset_name}/last': metrics})  # type: ignore[attr-defined]
-
-        # Save and push to hub
+            reward_model = None
+            try:
+                if explicit_policy_eval:
+                    reward_model = load_reward_model_for_task(
+                        training_args.reward_model_path,
+                        task_index,
+                        torch_dtype,
+                        model_args.trust_remote_code,
+                    )
+                    metrics.update(trainer.evaluate_policy(reward_model=reward_model))
+                    if training_args.log_completions:
+                        trainer.generate_completions_table(
+                            reward_model=reward_model,
+                            max_batches=training_args.completion_logging_batches,
+                        )
+            finally:
+                if reward_model is not None:
+                    del reward_model
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+
+            metrics['dataset'] = task_index
+            trainer.log_metrics(f'eval/dataset/{task_index}', metrics)
+            trainer.save_metrics('eval', metrics)
+            if wb.run is not None:
+                wb.log({'eval/last': metrics})
+                wb.log({f'task/{current_dataset_name}/last': metrics})
+
         trainer.save_model(os.path.join(training_args.output_dir, 'last'))
         if training_args.push_to_hub:
             trainer.push_to_hub(
-                dataset_name=(
-                    'Continual_DPO_' + script_args.dataset_name + '_' + str(i),
-                )
+                dataset_name='Continual_DPO_' + script_args.dataset_name + '_' + str(task_index)
             )
 
-        # If using DeepSpeed through Accelerate, tear down the engine after training.
-        if hasattr(trainer, 'deepspeed') and trainer.deepspeed is not None:
-            # Remove reference to the DeepSpeed engine to allow proper cleanup.
-            del trainer.deepspeed
-        # Free cached GPU memory.
-        torch.cuda.empty_cache()
-
 
 if __name__ == '__main__':
     dataclass_types = (ContinualDPOArguments, ContinualDPOConfig, ModelConfig)
diff --git a/benchmarks/dpo/sweep_configs/dpo_sweep.yaml b/benchmarks/dpo/sweep_configs/dpo_sweep.yaml
index 3af79cb1..a6866bc7 100644
--- a/benchmarks/dpo/sweep_configs/dpo_sweep.yaml
+++ b/benchmarks/dpo/sweep_configs/dpo_sweep.yaml
@@ -11,13 +11,13 @@ parameters:
   num_train_epochs:
     values: [1, 2, 3]
   per_device_train_batch_size:
-    values: [1, 2, 4]
+    values: [1, 2]
   gradient_checkpointing:
     values: [true, false]
   eval_steps:
-    values: [50, 100]
+    values: [100, 200]
 command:
-  - uv run
+  - python
   - benchmarks/dpo/dpo_continual.py
   - --dataset_name
   - debug
@@ -36,7 +36,15 @@ command:
   - --eval_steps
   - ${eval_steps}
   - --logging_steps
-  - 25
-  - --run_output_dir
+  - 100
+  - --bf16
+  - --output_dir
   - Qwen2-0.5B-DPO
+  - --max_prompt_length
+  - 256
+  - --max_completion_length
+  - 256
+  - --max_length
+  - 512
+  - --use_peft
   - --no_remove_unused_columns
diff --git a/test/test_benchmarks/test_dpo_benchmark_regressions.py b/test/test_benchmarks/test_dpo_benchmark_regressions.py
new file mode 100644
index 00000000..1e98e28a
--- /dev/null
+++ b/test/test_benchmarks/test_dpo_benchmark_regressions.py
@@ -0,0 +1,121 @@
+import ast
+from pathlib import Path
+
+REPO_ROOT = Path('/home/runner/work/AIF-Gen/AIF-Gen')
+TRAINER_PATH = REPO_ROOT / 'benchmarks/dpo/continual_dpo_trainer.py'
+SCRIPT_PATH = REPO_ROOT / 'benchmarks/dpo/dpo_continual.py'
+
+
+def parse_module(path: Path) -> ast.Module:
+    return ast.parse(path.read_text())
+
+
+def get_class(module: ast.Module, class_name: str) -> ast.ClassDef:
+    for node in module.body:
+        if isinstance(node, ast.ClassDef) and node.name == class_name:
+            return node
+    raise AssertionError(f'Class {class_name} not found in {module}')
+
+
+def get_function(module: ast.Module, function_name: str) -> ast.FunctionDef:
+    for node in module.body:
+        if isinstance(node, ast.FunctionDef) and node.name == function_name:
+            return node
+    raise AssertionError(f'Function {function_name} not found in {module}')
+
+
+def get_method(class_node: ast.ClassDef, method_name: str) -> ast.FunctionDef:
+    for node in class_node.body:
+        if isinstance(node, ast.FunctionDef) and node.name == method_name:
+            return node
+    raise AssertionError(f'Method {method_name} not found in {class_node.name}')
+
+
+def called_attribute_names(node: ast.AST) -> list[str]:
+    names = []
+    for call in ast.walk(node):
+        if isinstance(call, ast.Call) and isinstance(call.func, ast.Attribute):
+            names.append(call.func.attr)
+    return names
+
+
+def test_trainer_no_longer_overrides_generic_log_or_accelerator_creation() -> None:
+    trainer_module = parse_module(TRAINER_PATH)
+    trainer_class = get_class(trainer_module, 'ContinualDPOTrainer')
+    method_names = {
+        node.name for node in trainer_class.body if isinstance(node, ast.FunctionDef)
+    }
+
+    assert 'log' not in method_names
+    assert 'create_accelerator_and_postprocess' not in method_names
+    assert 'set_task_datasets' in method_names
+    assert 'generate_completions_table' in method_names
+
+
+
+def test_generate_completions_uses_prompt_tokens_for_generation_samples() -> None:
+    trainer_module = parse_module(TRAINER_PATH)
+    trainer_class = get_class(trainer_module, 'ContinualDPOTrainer')
+    method = get_method(trainer_class, 'generate_completions_table')
+
+    subscripts = []
+    for node in ast.walk(method):
+        if isinstance(node, ast.Subscript) and isinstance(node.value, ast.Name):
+            if node.value.id == 'batch' and isinstance(node.slice, ast.Constant):
+                subscripts.append(node.slice.value)
+
+    assert 'prompt_input_ids' in subscripts
+
+
+
+def test_dpo_script_uses_single_trainer_lifecycle_and_task_switch_method() -> None:
+    script_module = parse_module(SCRIPT_PATH)
+    main_function = get_function(script_module, 'main')
+
+    trainer_inits = [
+        node
+        for node in ast.walk(main_function)
+        if isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Name)
+        and node.func.id == 'ContinualDPOTrainer'
+    ]
+    assert len(trainer_inits) == 1
+
+    attr_calls = called_attribute_names(main_function)
+    assert 'set_task_datasets' in attr_calls
+    assert 'evaluate_policy' in attr_calls
+    assert 'generate_completions_table' in attr_calls
+
+    trainer_log_calls = [
+        node
+        for node in ast.walk(main_function)
+        if isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Attribute)
+        and node.func.attr == 'log'
+        and isinstance(node.func.value, ast.Name)
+        and node.func.value.id == 'trainer'
+    ]
+    assert trainer_log_calls == []
+
+
+
+def test_dpo_script_drops_duplicate_local_import_and_keeps_package_import() -> None:
+    script_module = parse_module(SCRIPT_PATH)
+    imported_modules = [
+        node.module
+        for node in script_module.body
+        if isinstance(node, ast.ImportFrom)
+    ]
+
+    assert 'continual_dpo_trainer' not in imported_modules
+    assert 'benchmarks.dpo.continual_dpo_trainer' in imported_modules
+
+
+
+def test_dpo_script_has_explicit_reward_model_loader_helper() -> None:
+    script_module = parse_module(SCRIPT_PATH)
+    load_helper = get_function(script_module, 'load_reward_model_for_task')
+    source = ast.unparse(load_helper)
+
+    assert 'AutoModelForSequenceClassification.from_pretrained' in source
+    assert 'torch_dtype' in source

From 57fbb4342e6ec3d5158fd119d4f98aad45df7c19 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 18:38:08 +0000
Subject: [PATCH 03/16] Address DPO review fixes

---
 benchmarks/dpo/sweep_configs/dpo_sweep.yaml            | 3 ++-
 test/test_benchmarks/test_dpo_benchmark_regressions.py | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/benchmarks/dpo/sweep_configs/dpo_sweep.yaml b/benchmarks/dpo/sweep_configs/dpo_sweep.yaml
index a6866bc7..4076b7c0 100644
--- a/benchmarks/dpo/sweep_configs/dpo_sweep.yaml
+++ b/benchmarks/dpo/sweep_configs/dpo_sweep.yaml
@@ -17,7 +17,8 @@ parameters:
   eval_steps:
     values: [100, 200]
 command:
-  - python
+  - uv
+  - run
   - benchmarks/dpo/dpo_continual.py
   - --dataset_name
   - debug
diff --git a/test/test_benchmarks/test_dpo_benchmark_regressions.py b/test/test_benchmarks/test_dpo_benchmark_regressions.py
index 1e98e28a..da07236e 100644
--- a/test/test_benchmarks/test_dpo_benchmark_regressions.py
+++ b/test/test_benchmarks/test_dpo_benchmark_regressions.py
@@ -1,7 +1,7 @@
 import ast
 from pathlib import Path
 
-REPO_ROOT = Path('/home/runner/work/AIF-Gen/AIF-Gen')
+REPO_ROOT = Path(__file__).resolve().parents[2]
 TRAINER_PATH = REPO_ROOT / 'benchmarks/dpo/continual_dpo_trainer.py'
 SCRIPT_PATH = REPO_ROOT / 'benchmarks/dpo/dpo_continual.py'
 

From 87f8a749b049f12338789a1b64e41f589fe54f9d Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 18:39:13 +0000
Subject: [PATCH 04/16] Polish continual DPO edge cases

---
 benchmarks/dpo/accelerate_configs/deepspeed_zero3.yaml | 2 +-
 benchmarks/dpo/continual_dpo_trainer.py                | 5 ++++-
 benchmarks/dpo/dpo_continual.py                        | 9 ++++++---
 3 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/benchmarks/dpo/accelerate_configs/deepspeed_zero3.yaml b/benchmarks/dpo/accelerate_configs/deepspeed_zero3.yaml
index 4b4312ea..cda38384 100644
--- a/benchmarks/dpo/accelerate_configs/deepspeed_zero3.yaml
+++ b/benchmarks/dpo/accelerate_configs/deepspeed_zero3.yaml
@@ -13,7 +13,7 @@ machine_rank: 0
 main_training_function: main
 mixed_precision: bf16
 num_machines: 1
-num_processes: 4  # Set this to the number of GPUs you are actually launching.
+num_processes: 1  # Placeholder: edit this to your actual GPU count before multi-GPU launch.
 rdzv_backend: static
 same_network: true
 tpu_env: []
diff --git a/benchmarks/dpo/continual_dpo_trainer.py b/benchmarks/dpo/continual_dpo_trainer.py
index f6dd8a28..026eadeb 100644
--- a/benchmarks/dpo/continual_dpo_trainer.py
+++ b/benchmarks/dpo/continual_dpo_trainer.py
@@ -272,7 +272,10 @@ def reward_model_context(
         disable_dropout_in_model(active_reward_model)
         active_reward_model.eval()
 
-        original_device = next(active_reward_model.parameters()).device
+        first_parameter = next(active_reward_model.parameters(), None)
+        original_device = (
+            first_parameter.device if first_parameter is not None else self.accelerator.device
+        )
         target_device = self.accelerator.device
         moved_to_accelerator = original_device != target_device
         if moved_to_accelerator:
diff --git a/benchmarks/dpo/dpo_continual.py b/benchmarks/dpo/dpo_continual.py
index fe72faf1..2cef4605 100644
--- a/benchmarks/dpo/dpo_continual.py
+++ b/benchmarks/dpo/dpo_continual.py
@@ -86,9 +86,12 @@ def validate_reward_model_paths(
                 reward_path,
                 num_labels=1,
             )
-        except Exception:
+        except Exception as exc:
             if not os.path.exists(reward_path):
-                raise ValueError(f'Reward model not found at {reward_path}') from None
+                raise ValueError(f'Reward model not found at {reward_path}') from exc
+            raise ValueError(
+                f'Failed to load reward model at {reward_path}: {exc}'
+            ) from exc
 
 
 def load_reward_model_for_task(
@@ -105,7 +108,7 @@ def load_reward_model_for_task(
         'num_labels': 1,
         'trust_remote_code': trust_remote_code,
     }
-    if isinstance(torch_dtype, torch.dtype):
+    if torch_dtype is not None:
         reward_model_kwargs['torch_dtype'] = torch_dtype
 
     return AutoModelForSequenceClassification.from_pretrained(

From c172f504e6d997db70dff7db26106e3a968c807b Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 18:40:02 +0000
Subject: [PATCH 05/16] Tidy DPO script imports

---
 benchmarks/dpo/dpo_continual.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/benchmarks/dpo/dpo_continual.py b/benchmarks/dpo/dpo_continual.py
index 2cef4605..88d324c9 100644
--- a/benchmarks/dpo/dpo_continual.py
+++ b/benchmarks/dpo/dpo_continual.py
@@ -21,6 +21,7 @@
 from trl.trainer.utils import SIMPLE_CHAT_TEMPLATE
 
 import wandb as wb
+
 from benchmarks.dataloading import init_continual_dataset
 from benchmarks.dpo.continual_dpo_trainer import (
     ContinualDPOArguments,

From 5371ccf3cf0ff4d7b271028ada07aedc0786ffb9 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 18:41:01 +0000
Subject: [PATCH 06/16] Clarify DPO eval dataloaders

---
 benchmarks/dpo/continual_dpo_trainer.py | 4 ++++
 benchmarks/dpo/dpo_continual.py         | 1 -
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/benchmarks/dpo/continual_dpo_trainer.py b/benchmarks/dpo/continual_dpo_trainer.py
index 026eadeb..0cf5cea7 100644
--- a/benchmarks/dpo/continual_dpo_trainer.py
+++ b/benchmarks/dpo/continual_dpo_trainer.py
@@ -325,6 +325,8 @@ def evaluate_policy(
 
             with torch.inference_mode():
                 for batch in self.eval_policy_dataloader:
+                    # `eval_policy_dataloader` is built from `preprocess_policy_dataset`, which stores prompts under
+                    # `input_ids` for reward-based policy evaluation only.
                     query = batch['input_ids'].to(self.accelerator.device)
                     context_length = query.shape[1]
                     with unwrap_model_for_generation(
@@ -394,6 +396,8 @@ def generate_completions_table(
                         if batch_index >= batches_to_log:
                             break
 
+                        # `get_eval_dataloader()` uses the tokenized DPO eval dataset, where prompts are kept under
+                        # `prompt_input_ids` together with the chosen/rejected preference targets.
                         query = batch['prompt_input_ids'].to(self.accelerator.device)
                         context_length = query.shape[1]
                         query_response, _ = batch_generation(
diff --git a/benchmarks/dpo/dpo_continual.py b/benchmarks/dpo/dpo_continual.py
index 88d324c9..62f6302c 100644
--- a/benchmarks/dpo/dpo_continual.py
+++ b/benchmarks/dpo/dpo_continual.py
@@ -81,7 +81,6 @@ def validate_reward_model_paths(
 
     for task_index in range(num_tasks):
         reward_path = get_task_reward_model_path(reward_model_root, task_index)
-        assert reward_path is not None
         try:
             AutoModelForSequenceClassification.from_pretrained(
                 reward_path,

From 74e1b31286b2ba30e8cf4e41d78d38f714dfb711 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 18:42:15 +0000
Subject: [PATCH 07/16] Refine DPO validation and docs

---
 benchmarks/dpo/README.md                      |  9 +++-----
 .../accelerate_configs/deepspeed_zero3.yaml   |  2 +-
 benchmarks/dpo/continual_dpo_trainer.py       |  6 +++++
 benchmarks/dpo/dpo_continual.py               | 23 +++++++++++--------
 4 files changed, 23 insertions(+), 17 deletions(-)

diff --git a/benchmarks/dpo/README.md b/benchmarks/dpo/README.md
index 08643f07..64b617b3 100644
--- a/benchmarks/dpo/README.md
+++ b/benchmarks/dpo/README.md
@@ -2,17 +2,14 @@
 
 ## Install benchmark dependencies
 
-```sh
-pip install -e .
-pip install pytest ruff accelerate==0.34.2 deepspeed==0.16.3 datasets>=3.2.0 numpy pandas wandb transformers>=4.49.0 trl>=0.15.2 peft>=0.14.0
-```
-
-If you use the project's preferred workflow, sync the benchmark group instead:
+The preferred setup is:
 
 ```sh
 uv sync --group benchmarks --group dev
 ```
 
+If you manage dependencies with plain pip instead, install the repository plus the benchmark extras listed in `pyproject.toml` (for example: `pytest`, `ruff`, `transformers`, `trl`, `accelerate`, `deepspeed`, `datasets`, `numpy`, `pandas`, `wandb`, and `peft`).
+
 ## What changed in this benchmark
 
 - The continual DPO pipeline now keeps **one trainer/model lifecycle** and swaps task datasets safely instead of reusing a shared `Accelerator` across multiple trainers.
diff --git a/benchmarks/dpo/accelerate_configs/deepspeed_zero3.yaml b/benchmarks/dpo/accelerate_configs/deepspeed_zero3.yaml
index cda38384..acf00796 100644
--- a/benchmarks/dpo/accelerate_configs/deepspeed_zero3.yaml
+++ b/benchmarks/dpo/accelerate_configs/deepspeed_zero3.yaml
@@ -13,7 +13,7 @@ machine_rank: 0
 main_training_function: main
 mixed_precision: bf16
 num_machines: 1
-num_processes: 1  # Placeholder: edit this to your actual GPU count before multi-GPU launch.
+num_processes: 1  # Placeholder: edit to your GPU count before multi-GPU launch (for example 2, 4, or 8).
 rdzv_backend: static
 same_network: true
 tpu_env: []
diff --git a/benchmarks/dpo/continual_dpo_trainer.py b/benchmarks/dpo/continual_dpo_trainer.py
index 0cf5cea7..aea8f285 100644
--- a/benchmarks/dpo/continual_dpo_trainer.py
+++ b/benchmarks/dpo/continual_dpo_trainer.py
@@ -224,6 +224,12 @@ def set_reward_model(
         self,
         reward_model: Optional[Union[PreTrainedModel, nn.Module, str]],
     ) -> None:
+        """Store a reward model for explicit evaluation-only flows.
+
+        This does not move or DeepSpeed-wrap the reward model during training; it only updates
+        the optional model used by ``reward_model_context`` when explicit policy evaluation or
+        completion logging is requested.
+        """
         self.reward_model = reward_model
         if self.reward_model is not None and not isinstance(self.reward_model, str):
             disable_dropout_in_model(self.reward_model)
diff --git a/benchmarks/dpo/dpo_continual.py b/benchmarks/dpo/dpo_continual.py
index 62f6302c..09ebbe37 100644
--- a/benchmarks/dpo/dpo_continual.py
+++ b/benchmarks/dpo/dpo_continual.py
@@ -7,6 +7,7 @@
 import torch
 from datasets import Dataset
 from transformers import (
+    AutoConfig,
     AutoModelForCausalLM,
     AutoModelForSequenceClassification,
     AutoTokenizer,
@@ -82,10 +83,7 @@ def validate_reward_model_paths(
     for task_index in range(num_tasks):
         reward_path = get_task_reward_model_path(reward_model_root, task_index)
         try:
-            AutoModelForSequenceClassification.from_pretrained(
-                reward_path,
-                num_labels=1,
-            )
+            AutoConfig.from_pretrained(reward_path, trust_remote_code=True)
         except Exception as exc:
             if not os.path.exists(reward_path):
                 raise ValueError(f'Reward model not found at {reward_path}') from exc
@@ -230,12 +228,17 @@ def main(
                         torch_dtype,
                         model_args.trust_remote_code,
                     )
-                    metrics.update(trainer.evaluate_policy(reward_model=reward_model))
-                    if training_args.log_completions:
-                        trainer.generate_completions_table(
-                            reward_model=reward_model,
-                            max_batches=training_args.completion_logging_batches,
-                        )
+                    try:
+                        metrics.update(trainer.evaluate_policy(reward_model=reward_model))
+                        if training_args.log_completions:
+                            trainer.generate_completions_table(
+                                reward_model=reward_model,
+                                max_batches=training_args.completion_logging_batches,
+                            )
+                    except Exception as exc:
+                        raise RuntimeError(
+                            f'Explicit policy evaluation failed for task {task_index} ({current_dataset_name}).'
+                        ) from exc
             finally:
                 if reward_model is not None:
                     del reward_model

From 1f3f46360d8255304d75209040f71685c6c77fba Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 18:43:07 +0000
Subject: [PATCH 08/16] Polish DPO script formatting

---
 benchmarks/dpo/dpo_continual.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/benchmarks/dpo/dpo_continual.py b/benchmarks/dpo/dpo_continual.py
index 09ebbe37..1372df9f 100644
--- a/benchmarks/dpo/dpo_continual.py
+++ b/benchmarks/dpo/dpo_continual.py
@@ -155,8 +155,7 @@ def main(
         ref_model = None
 
     tokenizer = AutoTokenizer.from_pretrained(
-        model_args.model_name_or_path,
-        trust_remote_code=model_args.trust_remote_code,
+        model_args.model_name_or_path, trust_remote_code=model_args.trust_remote_code
     )
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
@@ -255,7 +254,7 @@ def main(
         trainer.save_model(os.path.join(training_args.output_dir, 'last'))
         if training_args.push_to_hub:
             trainer.push_to_hub(
-                dataset_name='Continual_DPO_' + script_args.dataset_name + '_' + str(task_index)
+                dataset_name=f'Continual_DPO_{script_args.dataset_name}_{task_index}'
             )
 
 

From 767b6a8e405213e7c027ef87811c7d9d3b4277a6 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 19:03:07 +0000
Subject: [PATCH 09/16] Add opt-in profiling for continual DPO training

---
 benchmarks/dpo/continual_dpo_trainer.py |  73 +++++++-
 benchmarks/dpo/dpo_continual.py         | 215 +++++++++++++++++-------
 2 files changed, 229 insertions(+), 59 deletions(-)

diff --git a/benchmarks/dpo/continual_dpo_trainer.py b/benchmarks/dpo/continual_dpo_trainer.py
index aea8f285..1cd3619b 100644
--- a/benchmarks/dpo/continual_dpo_trainer.py
+++ b/benchmarks/dpo/continual_dpo_trainer.py
@@ -1,8 +1,9 @@
 import os
+import time
 from collections import defaultdict
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from typing import Callable, Optional, Union
+from typing import Any, Callable, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -112,6 +113,76 @@ class ContinualDPOConfig(DPOConfig):
             'help': 'Number of evaluation batches to sample when --log_completions is enabled.'
         },
     )
+    enable_profiling: bool = field(
+        default=False,
+        metadata={'help': 'Enable profiling hooks for continual DPO runs.'},
+    )
+    profiling_steps: int = field(
+        default=3,
+        metadata={
+            'help': 'Number of active steps to capture with torch profiler when profiling is enabled.'
+        },
+    )
+    profile_memory: bool = field(
+        default=False,
+        metadata={'help': 'Enable CUDA memory tracing when profiling is enabled.'},
+    )
+    profile_output_dir: str = field(
+        default='profiles/continual_dpo',
+        metadata={'help': 'Output directory for torch profiler TensorBoard traces.'},
+    )
+
+
+class StepProfilingCallback(TrainerCallback):
+    def __init__(
+        self,
+        profiler: Optional[Any] = None,
+        profile_memory: bool = False,
+    ) -> None:
+        self.profiler = profiler
+        self.profile_memory = profile_memory
+        self._step_start_time: Optional[float] = None
+        self._step_time_total: float = 0.0
+        self._steps: int = 0
+
+    @override
+    def on_step_begin(self, args, state, control, **kwargs):
+        self._step_start_time = time.perf_counter()
+        return control
+
+    @override
+    def on_step_end(self, args, state, control, **kwargs):
+        if self._step_start_time is None:
+            return control
+
+        elapsed = time.perf_counter() - self._step_start_time
+        self._step_time_total += elapsed
+        self._steps += 1
+
+        world_size = max(1, int(getattr(args, 'world_size', 1)))
+        global_batch_size = int(args.per_device_train_batch_size) * world_size
+        logs: dict[str, float] = {
+            'step': float(state.global_step),
+            'profiling/step_time_s': float(elapsed),
+            'profiling/step_time_avg_s': float(self._step_time_total / self._steps),
+            'profiling/samples_per_sec': float(global_batch_size / max(elapsed, 1e-12)),
+        }
+
+        if self.profile_memory and torch.cuda.is_available():
+            logs['profiling/gpu_memory_allocated_gb'] = float(
+                torch.cuda.memory_allocated() / (1024**3)
+            )
+            logs['profiling/gpu_memory_reserved_gb'] = float(
+                torch.cuda.memory_reserved() / (1024**3)
+            )
+
+        state.log_history.append(logs)
+
+        if self.profiler is not None:
+            self.profiler.step()
+
+        self._step_start_time = None
+        return control
 
 
 class ContinualDPOTrainer(DPOTrainer):
diff --git a/benchmarks/dpo/dpo_continual.py b/benchmarks/dpo/dpo_continual.py
index 1372df9f..268ae056 100644
--- a/benchmarks/dpo/dpo_continual.py
+++ b/benchmarks/dpo/dpo_continual.py
@@ -1,7 +1,9 @@
 """Adaptation of the DPO TRL training script for continual learning."""
 
 import os
+import time
 import warnings
+from contextlib import contextmanager
 from typing import Optional
 
 import torch
@@ -28,6 +30,7 @@
     ContinualDPOArguments,
     ContinualDPOConfig,
     ContinualDPOTrainer,
+    StepProfilingCallback,
 )
 
 
@@ -115,6 +118,44 @@ def load_reward_model_for_task(
     )
 
 
+def _build_torch_profiler(training_args: ContinualDPOConfig):
+    if not training_args.enable_profiling:
+        return None
+
+    active_steps = max(1, int(training_args.profiling_steps))
+    os.makedirs(training_args.profile_output_dir, exist_ok=True)
+
+    activities = [torch.profiler.ProfilerActivity.CPU]
+    if torch.cuda.is_available():
+        activities.append(torch.profiler.ProfilerActivity.CUDA)
+
+    return torch.profiler.profile(
+        activities=activities,
+        schedule=torch.profiler.schedule(wait=1, warmup=1, active=active_steps),
+        on_trace_ready=torch.profiler.tensorboard_trace_handler(
+            training_args.profile_output_dir
+        ),
+        record_shapes=True,
+        profile_memory=training_args.profile_memory,
+    )
+
+
+@contextmanager
+def _time_phase(label: str, wandb_run, enabled: bool):
+    if not enabled:
+        yield
+        return
+
+    start = time.perf_counter()
+    try:
+        yield
+    finally:
+        elapsed = time.perf_counter() - start
+        print(f'[{label}] elapsed: {elapsed:.3f}s')
+        if wandb_run is not None:
+            wb.log({f'profiling/{label}_time_s': float(elapsed)})
+
+
 def main(
     script_args: ContinualDPOArguments,
     training_args: ContinualDPOConfig,
@@ -167,12 +208,13 @@ def main(
             name for name, buffer in model.named_buffers() if buffer.dtype == torch.bool
         ]
 
-    continual_dataset: list[dict[str, Dataset]] = init_continual_dataset(
-        script_args.dataset_name,
-        mock=training_args.mock,
-        tokenizer=tokenizer,
-        tools=training_args.tools,
-    )
+    with _time_phase('data_loading', wb.run, training_args.enable_profiling):
+        continual_dataset: list[dict[str, Dataset]] = init_continual_dataset(
+            script_args.dataset_name,
+            mock=training_args.mock,
+            tokenizer=tokenizer,
+            tools=training_args.tools,
+        )
     output_dir = training_args.output_dir
     eval_enabled = training_args.eval_strategy != 'no'
     explicit_policy_eval = training_args.eval_policy_metrics or training_args.log_completions
@@ -196,66 +238,123 @@ def main(
         eval_dataset=first_dataset.get(script_args.dataset_test_split),
         peft_config=peft_config,
     )
+    profiler = _build_torch_profiler(training_args)
+    if training_args.enable_profiling:
+        trainer.add_callback(
+            StepProfilingCallback(
+                profiler=profiler,
+                profile_memory=training_args.profile_memory,
+            )
+        )
 
     if wb.run is not None:
         wb.log({'dataset/name': script_args.dataset_name})
+    if training_args.enable_profiling and torch.cuda.is_available():
+        print('CUDA memory summary before continual loop:')
+        print(torch.cuda.memory_summary())
 
+    first_task_profiled = False
     for task_index, dataset in enumerate(continual_dataset):
-        current_dataset_name = f'dataset-{task_index}'
-        training_args.output_dir = f'{output_dir}/{current_dataset_name}'
-        trainer.args.output_dir = training_args.output_dir
-
-        if task_index > 0:
-            trainer.set_task_datasets(
-                train_dataset=dataset[script_args.dataset_train_split],
-                eval_dataset=dataset.get(script_args.dataset_test_split),
-                dataset_name=current_dataset_name,
+        with _time_phase(
+            f'task_{task_index}',
+            wb.run,
+            training_args.enable_profiling,
+        ):
+            current_dataset_name = f'dataset-{task_index}'
+            training_args.output_dir = f'{output_dir}/{current_dataset_name}'
+            trainer.args.output_dir = training_args.output_dir
+
+            if task_index > 0:
+                trainer.set_task_datasets(
+                    train_dataset=dataset[script_args.dataset_train_split],
+                    eval_dataset=dataset.get(script_args.dataset_test_split),
+                    dataset_name=current_dataset_name,
+                )
+
+            if training_args.enable_profiling and torch.cuda.is_available():
+                torch.cuda.reset_peak_memory_stats()
+                print(
+                    f'CUDA memory summary at task {task_index} start ({current_dataset_name}):'
+                )
+                print(torch.cuda.memory_summary())
+
+            print('Training dataset:', current_dataset_name)
+            should_profile_first_task = (
+                training_args.enable_profiling
+                and profiler is not None
+                and not first_task_profiled
             )
-
-        print('Training dataset:', current_dataset_name)
-        trainer.train()
-
-        should_run_task_eval = eval_enabled or explicit_policy_eval
-        if should_run_task_eval and trainer.eval_dataset is not None:
-            metrics = trainer.evaluate()
-            reward_model = None
+            if should_profile_first_task:
+                profiler.start()
             try:
-                if explicit_policy_eval:
-                    reward_model = load_reward_model_for_task(
-                        training_args.reward_model_path,
-                        task_index,
-                        torch_dtype,
-                        model_args.trust_remote_code,
-                    )
-                    try:
-                        metrics.update(trainer.evaluate_policy(reward_model=reward_model))
-                        if training_args.log_completions:
-                            trainer.generate_completions_table(
-                                reward_model=reward_model,
-                                max_batches=training_args.completion_logging_batches,
-                            )
-                    except Exception as exc:
-                        raise RuntimeError(
-                            f'Explicit policy evaluation failed for task {task_index} ({current_dataset_name}).'
-                        ) from exc
+                trainer.train()
             finally:
-                if reward_model is not None:
-                    del reward_model
-                if torch.cuda.is_available():
-                    torch.cuda.empty_cache()
-
-            metrics['dataset'] = task_index
-            trainer.log_metrics(f'eval/dataset/{task_index}', metrics)
-            trainer.save_metrics('eval', metrics)
-            if wb.run is not None:
-                wb.log({'eval/last': metrics})
-                wb.log({f'task/{current_dataset_name}/last': metrics})
-
-        trainer.save_model(os.path.join(training_args.output_dir, 'last'))
-        if training_args.push_to_hub:
-            trainer.push_to_hub(
-                dataset_name=f'Continual_DPO_{script_args.dataset_name}_{task_index}'
-            )
+                if should_profile_first_task:
+                    profiler.stop()
+                    first_task_profiled = True
+                    print(
+                        f'Profiler trace exported to {training_args.profile_output_dir} for task {task_index}.'
+                    )
+
+            should_run_task_eval = eval_enabled or explicit_policy_eval
+            if should_run_task_eval and trainer.eval_dataset is not None:
+                metrics = trainer.evaluate()
+                reward_model = None
+                try:
+                    if explicit_policy_eval:
+                        reward_model = load_reward_model_for_task(
+                            training_args.reward_model_path,
+                            task_index,
+                            torch_dtype,
+                            model_args.trust_remote_code,
+                        )
+                        try:
+                            metrics.update(trainer.evaluate_policy(reward_model=reward_model))
+                            if training_args.log_completions:
+                                trainer.generate_completions_table(
+                                    reward_model=reward_model,
+                                    max_batches=training_args.completion_logging_batches,
+                                )
+                        except Exception as exc:
+                            raise RuntimeError(
+                                f'Explicit policy evaluation failed for task {task_index} ({current_dataset_name}).'
+                            ) from exc
+                finally:
+                    if reward_model is not None:
+                        del reward_model
+                    if torch.cuda.is_available():
+                        torch.cuda.empty_cache()
+
+                metrics['dataset'] = task_index
+                trainer.log_metrics(f'eval/dataset/{task_index}', metrics)
+                trainer.save_metrics('eval', metrics)
+                if wb.run is not None:
+                    wb.log({'eval/last': metrics})
+                    wb.log({f'task/{current_dataset_name}/last': metrics})
+
+            if training_args.enable_profiling and torch.cuda.is_available():
+                peak_allocated_gb = float(torch.cuda.max_memory_allocated() / (1024**3))
+                peak_reserved_gb = float(torch.cuda.max_memory_reserved() / (1024**3))
+                print(
+                    f'Task {task_index} peak CUDA memory (GB): allocated={peak_allocated_gb:.3f}, reserved={peak_reserved_gb:.3f}'
+                )
+                if wb.run is not None:
+                    wb.log(
+                        {
+                            f'profiling/task_{task_index}_peak_memory_allocated_gb': peak_allocated_gb,
+                            f'profiling/task_{task_index}_peak_memory_reserved_gb': peak_reserved_gb,
+                        }
+                    )
+
+            trainer.save_model(os.path.join(training_args.output_dir, 'last'))
+            if training_args.push_to_hub:
+                trainer.push_to_hub(
+                    dataset_name=f'Continual_DPO_{script_args.dataset_name}_{task_index}'
+                )
+
+    if training_args.enable_profiling and torch.cuda.is_available():
+        print('CUDA memory summary after continual loop:')
+        print(torch.cuda.memory_summary())
 
 
 if __name__ == '__main__':

From bee90fe001c15ce6cd16e8bd2f54d1b5d3eec257 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 19:04:19 +0000
Subject: [PATCH 10/16] Refine continual DPO profiling instrumentation

---
 benchmarks/dpo/continual_dpo_trainer.py | 8 ++++----
 benchmarks/dpo/dpo_continual.py         | 6 ++++--
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/benchmarks/dpo/continual_dpo_trainer.py b/benchmarks/dpo/continual_dpo_trainer.py
index 1cd3619b..0d5eb804 100644
--- a/benchmarks/dpo/continual_dpo_trainer.py
+++ b/benchmarks/dpo/continual_dpo_trainer.py
@@ -159,10 +159,10 @@ def on_step_end(self, args, state, control, **kwargs):
         self._step_time_total += elapsed
         self._steps += 1
 
-        world_size = max(1, int(getattr(args, 'world_size', 1)))
-        global_batch_size = int(args.per_device_train_batch_size) * world_size
-        logs: dict[str, float] = {
-            'step': float(state.global_step),
+        world_size = max(1, getattr(args, 'world_size', 1))
+        global_batch_size = args.per_device_train_batch_size * world_size
+        logs: dict[str, float | int] = {
+            'step': state.global_step,
             'profiling/step_time_s': float(elapsed),
             'profiling/step_time_avg_s': float(self._step_time_total / self._steps),
             'profiling/samples_per_sec': float(global_batch_size / max(elapsed, 1e-12)),
diff --git a/benchmarks/dpo/dpo_continual.py b/benchmarks/dpo/dpo_continual.py
index 268ae056..f3926924 100644
--- a/benchmarks/dpo/dpo_continual.py
+++ b/benchmarks/dpo/dpo_continual.py
@@ -122,7 +122,7 @@ def _build_torch_profiler(training_args: ContinualDPOConfig):
     if not training_args.enable_profiling:
         return None
 
-    active_steps = max(1, int(training_args.profiling_steps))
+    active_steps = max(1, training_args.profiling_steps)
     os.makedirs(training_args.profile_output_dir, exist_ok=True)
 
     activities = [torch.profiler.ProfilerActivity.CPU]
@@ -220,7 +220,9 @@ def main(
     explicit_policy_eval = training_args.eval_policy_metrics or training_args.log_completions
 
     if training_args.eval_policy_metrics and training_args.reward_model_path is None:
-        raise ValueError('--eval_policy_metrics requires --reward_model_path.')
+        raise ValueError(
+            'Cannot use --eval_policy_metrics without --reward_model_path; reward model path must be specified for policy evaluation.'
+        )
 
     if explicit_policy_eval:
         validate_reward_model_paths(

From 91cdcf97f24043c2a9e21b8c33e6fb69262a8613 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 19:05:08 +0000
Subject: [PATCH 11/16] Polish profiling typings and docs

---
 benchmarks/dpo/continual_dpo_trainer.py | 2 ++
 benchmarks/dpo/dpo_continual.py         | 6 ++++--
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/benchmarks/dpo/continual_dpo_trainer.py b/benchmarks/dpo/continual_dpo_trainer.py
index 0d5eb804..ea52d1c9 100644
--- a/benchmarks/dpo/continual_dpo_trainer.py
+++ b/benchmarks/dpo/continual_dpo_trainer.py
@@ -134,6 +134,8 @@ class ContinualDPOConfig(DPOConfig):
 
 
 class StepProfilingCallback(TrainerCallback):
+    """Track per-step timing/memory metrics and optionally advance a torch profiler."""
+
     def __init__(
         self,
         profiler: Optional[Any] = None,
diff --git a/benchmarks/dpo/dpo_continual.py b/benchmarks/dpo/dpo_continual.py
index f3926924..9330a54f 100644
--- a/benchmarks/dpo/dpo_continual.py
+++ b/benchmarks/dpo/dpo_continual.py
@@ -118,7 +118,9 @@ def load_reward_model_for_task(
     )
 
 
-def _build_torch_profiler(training_args: ContinualDPOConfig):
+def _build_torch_profiler(
+    training_args: ContinualDPOConfig,
+) -> Optional[torch.profiler.profile]:
     if not training_args.enable_profiling:
         return None
 
@@ -141,7 +143,7 @@ def _build_torch_profiler(training_args: ContinualDPOConfig):
 
 
 @contextmanager
-def _time_phase(label: str, wandb_run, enabled: bool):
+def _time_phase(label: str, wandb_run: Optional[object], enabled: bool):
     if not enabled:
         yield
         return

From 4acb486a9c4aaffe56a9df999a2219170cd7db24 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 19:06:06 +0000
Subject: [PATCH 12/16] Harden profiling callback state handling

---
 benchmarks/dpo/continual_dpo_trainer.py | 8 ++++----
 benchmarks/dpo/dpo_continual.py         | 5 ++---
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/benchmarks/dpo/continual_dpo_trainer.py b/benchmarks/dpo/continual_dpo_trainer.py
index ea52d1c9..4fb14e52 100644
--- a/benchmarks/dpo/continual_dpo_trainer.py
+++ b/benchmarks/dpo/continual_dpo_trainer.py
@@ -154,10 +154,12 @@ def on_step_begin(self, args, state, control, **kwargs):
 
     @override
     def on_step_end(self, args, state, control, **kwargs):
-        if self._step_start_time is None:
+        step_start_time = self._step_start_time
+        self._step_start_time = None
+        if step_start_time is None:
             return control
 
-        elapsed = time.perf_counter() - self._step_start_time
+        elapsed = time.perf_counter() - step_start_time
         self._step_time_total += elapsed
         self._steps += 1
 
@@ -182,8 +184,6 @@ def on_step_end(self, args, state, control, **kwargs):
 
         if self.profiler is not None:
             self.profiler.step()
-
-        self._step_start_time = None
         return control
 
 
diff --git a/benchmarks/dpo/dpo_continual.py b/benchmarks/dpo/dpo_continual.py
index 9330a54f..8e546e18 100644
--- a/benchmarks/dpo/dpo_continual.py
+++ b/benchmarks/dpo/dpo_continual.py
@@ -4,7 +4,7 @@
 import time
 import warnings
 from contextlib import contextmanager
-from typing import Optional
+from typing import Any, Optional
 
 import torch
 from datasets import Dataset
@@ -143,7 +143,7 @@ def _build_torch_profiler(
 
 
 @contextmanager
-def _time_phase(label: str, wandb_run: Optional[object], enabled: bool):
+def _time_phase(label: str, wandb_run: Optional[Any], enabled: bool):
     if not enabled:
         yield
         return
@@ -285,7 +285,6 @@ def main(
             print('Training dataset:', current_dataset_name)
             should_profile_first_task = (
                 training_args.enable_profiling
-                and profiler is not None
                 and not first_task_profiled
             )
             if should_profile_first_task:

From 89a635c78285b813081207fbf7b53dd86e43531a Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 19:07:05 +0000
Subject: [PATCH 13/16] Document profiling helper APIs

---
 benchmarks/dpo/continual_dpo_trainer.py | 6 ++++++
 benchmarks/dpo/dpo_continual.py         | 6 ++++++
 2 files changed, 12 insertions(+)

diff --git a/benchmarks/dpo/continual_dpo_trainer.py b/benchmarks/dpo/continual_dpo_trainer.py
index 4fb14e52..b41cdefc 100644
--- a/benchmarks/dpo/continual_dpo_trainer.py
+++ b/benchmarks/dpo/continual_dpo_trainer.py
@@ -141,6 +141,12 @@ def __init__(
         profiler: Optional[Any] = None,
         profile_memory: bool = False,
     ) -> None:
+        """Initialize optional step-level profiler and memory metric tracking.
+
+        Args:
+            profiler: Torch profiler instance to advance once per training step.
+            profile_memory: Whether to collect CUDA allocated/reserved memory metrics.
+        """
         self.profiler = profiler
         self.profile_memory = profile_memory
         self._step_start_time: Optional[float] = None
diff --git a/benchmarks/dpo/dpo_continual.py b/benchmarks/dpo/dpo_continual.py
index 8e546e18..d2cd4192 100644
--- a/benchmarks/dpo/dpo_continual.py
+++ b/benchmarks/dpo/dpo_continual.py
@@ -40,6 +40,7 @@ def warn_for_memory_settings(
     training_args: ContinualDPOConfig,
     model_args: ModelConfig,
 ) -> None:
+    """Warn when configuration choices are likely to increase memory usage."""
     if training_args.max_completion_length is None:
         warnings.warn(
             'max_completion_length is unset. Long chosen/rejected responses can still cause large padded DPO batches; '
@@ -71,6 +72,7 @@ def get_task_reward_model_path(
     reward_model_root: Optional[str],
     task_index: int,
 ) -> Optional[str]:
+    """Build the task-specific reward-model path from the root path and task index."""
     if reward_model_root is None:
         return None
     return f'{reward_model_root}_{task_index}'
@@ -80,6 +82,7 @@ def validate_reward_model_paths(
     reward_model_root: Optional[str],
     num_tasks: int,
 ) -> None:
+    """Validate that each per-task reward model path exists and is loadable."""
     if reward_model_root is None:
         return
 
@@ -101,6 +104,7 @@ def load_reward_model_for_task(
     torch_dtype: Optional[torch.dtype],
     trust_remote_code: bool,
 ) -> Optional[AutoModelForSequenceClassification]:
+    """Load the reward model checkpoint associated with a specific continual task."""
     reward_path = get_task_reward_model_path(reward_model_root, task_index)
     if reward_path is None:
         return None
@@ -121,6 +125,7 @@ def load_reward_model_for_task(
 def _build_torch_profiler(
     training_args: ContinualDPOConfig,
 ) -> Optional[torch.profiler.profile]:
+    """Create a configured torch profiler when profiling is enabled, otherwise return None."""
     if not training_args.enable_profiling:
         return None
 
@@ -144,6 +149,7 @@ def _build_torch_profiler(
 
 @contextmanager
 def _time_phase(label: str, wandb_run: Optional[Any], enabled: bool):
+    """Time a logical phase and optionally log elapsed seconds to Weights & Biases."""
     if not enabled:
         yield
         return

From 45336103a3da1284ee3d98ab39873aca970510ff Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 19:07:55 +0000
Subject: [PATCH 14/16] Clarify continual DPO profiling error paths

---
 benchmarks/dpo/dpo_continual.py | 26 ++++++++++++++++----------
 1 file changed, 16 insertions(+), 10 deletions(-)

diff --git a/benchmarks/dpo/dpo_continual.py b/benchmarks/dpo/dpo_continual.py
index d2cd4192..1336bfca 100644
--- a/benchmarks/dpo/dpo_continual.py
+++ b/benchmarks/dpo/dpo_continual.py
@@ -228,9 +228,7 @@ def main(
     explicit_policy_eval = training_args.eval_policy_metrics or training_args.log_completions
 
     if training_args.eval_policy_metrics and training_args.reward_model_path is None:
-        raise ValueError(
-            'Cannot use --eval_policy_metrics without --reward_model_path; reward model path must be specified for policy evaluation.'
-        )
+        raise ValueError('Cannot use --eval_policy_metrics without --reward_model_path.')
 
     if explicit_policy_eval:
         validate_reward_model_paths(
@@ -317,17 +315,25 @@ def main(
                             torch_dtype,
                             model_args.trust_remote_code,
                         )
-                        try:
-                            metrics.update(trainer.evaluate_policy(reward_model=reward_model))
-                            if training_args.log_completions:
+                        if training_args.eval_policy_metrics:
+                            try:
+                                metrics.update(
+                                    trainer.evaluate_policy(reward_model=reward_model)
+                                )
+                            except Exception as exc:
+                                raise RuntimeError(
+                                    f'Reward-model policy evaluation failed for task {task_index} ({current_dataset_name}).'
+                                ) from exc
+                        if training_args.log_completions:
+                            try:
                                 trainer.generate_completions_table(
                                     reward_model=reward_model,
                                     max_batches=training_args.completion_logging_batches,
                                 )
-                        except Exception as exc:
-                            raise RuntimeError(
-                                f'Explicit policy evaluation failed for task {task_index} ({current_dataset_name}).'
-                            ) from exc
+                            except Exception as exc:
+                                raise RuntimeError(
+                                    f'Completion generation logging failed for task {task_index} ({current_dataset_name}).'
+                                ) from exc
                 finally:
                     if reward_model is not None:
                         del reward_model

From bfefda325fb8386a43cdf6bd58cdfe0ec49b00f8 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 19:08:49 +0000
Subject: [PATCH 15/16] Add profiling callback and timing docs

---
 benchmarks/dpo/continual_dpo_trainer.py | 2 ++
 benchmarks/dpo/dpo_continual.py         | 7 ++++++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/benchmarks/dpo/continual_dpo_trainer.py b/benchmarks/dpo/continual_dpo_trainer.py
index b41cdefc..173ee38c 100644
--- a/benchmarks/dpo/continual_dpo_trainer.py
+++ b/benchmarks/dpo/continual_dpo_trainer.py
@@ -155,11 +155,13 @@ def __init__(
 
     @override
     def on_step_begin(self, args, state, control, **kwargs):
+        """Capture step start timestamp before each optimizer step."""
         self._step_start_time = time.perf_counter()
         return control
 
     @override
     def on_step_end(self, args, state, control, **kwargs):
+        """Log per-step timing/memory metrics and advance profiler when available."""
         step_start_time = self._step_start_time
         self._step_start_time = None
         if step_start_time is None:
diff --git a/benchmarks/dpo/dpo_continual.py b/benchmarks/dpo/dpo_continual.py
index 1336bfca..0a38d973 100644
--- a/benchmarks/dpo/dpo_continual.py
+++ b/benchmarks/dpo/dpo_continual.py
@@ -149,7 +149,10 @@ def _build_torch_profiler(
 
 @contextmanager
 def _time_phase(label: str, wandb_run: Optional[Any], enabled: bool):
-    """Time a logical phase and optionally log elapsed seconds to Weights & Biases."""
+    """Time a logical phase and optionally log elapsed seconds to Weights & Biases.
+
+    When ``enabled`` is False this context manager is a no-op and does not collect timing.
+    """
     if not enabled:
         yield
         return
@@ -291,6 +294,8 @@ def main(
                 training_args.enable_profiling
                 and not first_task_profiled
             )
+            # Profile only the first task to capture representative kernels without duplicating
+            # large trace files across all continual tasks.
             if should_profile_first_task:
                 profiler.start()
             try:

From b17782a5e32ef9bfb2c948910184c0694812e07e Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Thu, 9 Jul 2026 19:20:02 +0000
Subject: [PATCH 16/16] Add memory profiling regression tests

---
 .../test_dpo_benchmark_regressions.py         | 34 +++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/test/test_benchmarks/test_dpo_benchmark_regressions.py b/test/test_benchmarks/test_dpo_benchmark_regressions.py
index da07236e..a92862f3 100644
--- a/test/test_benchmarks/test_dpo_benchmark_regressions.py
+++ b/test/test_benchmarks/test_dpo_benchmark_regressions.py
@@ -119,3 +119,37 @@ def test_dpo_script_has_explicit_reward_model_loader_helper() -> None:
 
     assert 'AutoModelForSequenceClassification.from_pretrained' in source
     assert 'torch_dtype' in source
+
+
+def test_step_profiling_callback_records_gpu_memory_metrics_when_enabled() -> None:
+    trainer_module = parse_module(TRAINER_PATH)
+    trainer_class = get_class(trainer_module, 'StepProfilingCallback')
+    method = get_method(trainer_class, 'on_step_end')
+    source = ast.unparse(method)
+
+    assert 'self.profile_memory and torch.cuda.is_available()' in source
+    assert 'torch.cuda.memory_allocated()' in source
+    assert 'torch.cuda.memory_reserved()' in source
+    assert 'profiling/gpu_memory_allocated_gb' in source
+    assert 'profiling/gpu_memory_reserved_gb' in source
+
+
+def test_torch_profiler_helper_threads_profile_memory_setting() -> None:
+    script_module = parse_module(SCRIPT_PATH)
+    profiler_helper = get_function(script_module, '_build_torch_profiler')
+    source = ast.unparse(profiler_helper)
+
+    assert 'profile_memory=training_args.profile_memory' in source
+
+
+def test_dpo_script_tracks_peak_cuda_memory_per_task_when_profiling() -> None:
+    script_module = parse_module(SCRIPT_PATH)
+    main_function = get_function(script_module, 'main')
+    source = ast.unparse(main_function)
+
+    assert 'torch.cuda.reset_peak_memory_stats()' in source
+    assert 'torch.cuda.max_memory_allocated()' in source
+    assert 'torch.cuda.max_memory_reserved()' in source
+    assert 'profiling/task_' in source
+    assert 'peak_memory_allocated_gb' in source
+    assert 'peak_memory_reserved_gb' in source