diff --git a/benchmarks/dpo/README.md b/benchmarks/dpo/README.md index 53cf7a48..64b617b3 100644 --- a/benchmarks/dpo/README.md +++ b/benchmarks/dpo/README.md @@ -1,92 +1,203 @@ -# Adaptation of TRL for Continual Learning +# Continual DPO benchmark -### Sync additional dependencies +## Install benchmark dependencies + +The preferred setup is: ```sh -uv sync --group benchmarks +uv sync --group benchmarks --group dev ``` -## Run DPO +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. +- 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. -### Lora +## 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 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 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 1 \ + --gradient_accumulation_steps 16 \ + --gradient_checkpointing \ + --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 ``` - -- 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..acf00796 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: 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/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..173ee38c 100644 --- a/benchmarks/dpo/continual_dpo_trainer.py +++ b/benchmarks/dpo/continual_dpo_trainer.py @@ -1,19 +1,20 @@ -import functools -import inspect 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 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 +29,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 +82,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 +95,107 @@ 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.' + }, + ) + 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 ContinualDPOTrainer(DPOTrainer): - # Shared accelerator instance across all trainer instances - shared_accelerator: Optional[Accelerator] = None - accelerator: Accelerator # now non-optional after creation +class StepProfilingCallback(TrainerCallback): + """Track per-step timing/memory metrics and optionally advance a torch profiler.""" + + def __init__( + self, + 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 + self._step_time_total: float = 0.0 + self._steps: int = 0 + @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: + return control + + elapsed = time.perf_counter() - step_start_time + self._step_time_total += elapsed + self._steps += 1 + + 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)), + } + + 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() + return control + + +class ContinualDPOTrainer(DPOTrainer): @override def __init__( self, @@ -156,83 +246,77 @@ 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: + """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) + 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 +334,59 @@ 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() + + 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: + 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 +397,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,9 +405,15 @@ 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: + # `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( @@ -301,12 +429,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 +439,101 @@ 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() + + 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 - queries = gather_object( - self.processing_class.batch_decode( - query, skip_special_tokens=True + # `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( + 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..0a38d973 100644 --- a/benchmarks/dpo/dpo_continual.py +++ b/benchmarks/dpo/dpo_continual.py @@ -1,15 +1,15 @@ """Adaptation of the DPO TRL training script for continual learning.""" import os +import time +import warnings +from contextlib import contextmanager +from typing import Any, Optional import torch -from continual_dpo_trainer import ( - ContinualDPOArguments, - ContinualDPOConfig, - ContinualDPOTrainer, -) from datasets import Dataset from transformers import ( + AutoConfig, AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer, @@ -24,21 +24,154 @@ 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, ContinualDPOConfig, ContinualDPOTrainer, + StepProfilingCallback, ) # 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: + """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; ' + '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]: + """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}' + + +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 + + for task_index in range(num_tasks): + reward_path = get_task_reward_model_path(reward_model_root, task_index) + try: + 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 + raise ValueError( + f'Failed to load reward model at {reward_path}: {exc}' + ) from exc + + +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]: + """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 + + reward_model_kwargs = { + 'num_labels': 1, + 'trust_remote_code': trust_remote_code, + } + if torch_dtype is not None: + reward_model_kwargs['torch_dtype'] = torch_dtype + + return AutoModelForSequenceClassification.from_pretrained( + reward_path, + **reward_model_kwargs, + ) + + +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 + + active_steps = max(1, 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: Optional[Any], enabled: bool): + """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 + + 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, 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 +180,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,7 +206,6 @@ 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 ) @@ -82,94 +214,167 @@ def main( 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, - 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 - # 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}' - training_args.output_dir = f'{output_dir}/{current_dataset_name}' - - # 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 training_args.eval_policy_metrics and training_args.reward_model_path is None: + raise ValueError('Cannot use --eval_policy_metrics without --reward_model_path.') - 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, + if explicit_policy_eval: + validate_reward_model_paths( + training_args.reward_model_path, + len(continual_dataset), ) - # 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': - 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 - 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), + 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, + ) + 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): + 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 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: + trainer.train() + finally: + 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, + ) + 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'Completion generation logging 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 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 training_args.enable_profiling and torch.cuda.is_available(): + print('CUDA memory summary after continual loop:') + print(torch.cuda.memory_summary()) if __name__ == '__main__': diff --git a/benchmarks/dpo/sweep_configs/dpo_sweep.yaml b/benchmarks/dpo/sweep_configs/dpo_sweep.yaml index 3af79cb1..4076b7c0 100644 --- a/benchmarks/dpo/sweep_configs/dpo_sweep.yaml +++ b/benchmarks/dpo/sweep_configs/dpo_sweep.yaml @@ -11,13 +11,14 @@ 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 + - uv + - run - benchmarks/dpo/dpo_continual.py - --dataset_name - debug @@ -36,7 +37,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..a92862f3 --- /dev/null +++ b/test/test_benchmarks/test_dpo_benchmark_regressions.py @@ -0,0 +1,155 @@ +import ast +from pathlib import Path + +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' + + +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 + + +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