diff --git a/docs/.vitepress/config.mts b/docs/.vitepress/config.mts index c2157f054..7feecf001 100644 --- a/docs/.vitepress/config.mts +++ b/docs/.vitepress/config.mts @@ -244,6 +244,7 @@ export default defineConfig({ { text: 'Quick Start', link: '/en/guide/quick-start' }, { text: 'Customize Training', link: '/en/guide/customize-training' }, { text: 'SFT Training', link: '/en/guide/sft-training' }, + { text: 'DPO Training', link: '/en/guide/dpo-training' }, { text: 'PPO Training', link: '/en/guide/ppo-training' }, { text: 'REINFORCE++', link: '/en/guide/reinforce-plus-plus' }, { text: 'REINFORCE++ Report', link: '/en/guide/reinforce-plus-plus-training-report' }, @@ -356,6 +357,7 @@ export default defineConfig({ { text: '快速上手', link: '/zh/guide/quick-start' }, { text: '自定义训练', link: '/zh/guide/customize-training' }, { text: 'SFT 训练', link: '/zh/guide/sft-training' }, + { text: 'DPO 训练', link: '/zh/guide/dpo-training' }, { text: 'PPO 训练', link: '/zh/guide/ppo-training' }, { text: 'REINFORCE++', link: '/zh/guide/reinforce-plus-plus' }, { text: 'REINFORCE++ 训练与数值验证报告', link: '/zh/guide/reinforce-plus-plus-training-report' }, diff --git a/docs/en/guide/dpo-training.md b/docs/en/guide/dpo-training.md new file mode 100644 index 000000000..29a1af12a --- /dev/null +++ b/docs/en/guide/dpo-training.md @@ -0,0 +1,64 @@ +# DPO Training + +Relax supports Direct Preference Optimization (DPO) through the offline SFT data path. The public Task 31 recipe is [`run-qwen3-0.6B-ultrafeedback-1xgpu.sh`](../../../scripts/training/dpo/run-qwen3-0.6B-ultrafeedback-1xgpu.sh). + +## Prepare the preference subset + +Generate the deterministic UltraFeedback subset from its pinned dataset revision: + +```bash +python scripts/data/prepare_ultrafeedback_preferences.py \ + --output-dir /data/task31-ultrafeedback +``` + +The command creates train/eval JSONL and Parquet files plus `manifest.json`. For the published Task 31 subset, compare the generated manifest with the [reproducibility evidence bundle](https://github.com/user-attachments/files/31305744/task31-pr1-dpo-evidence-public-v2.tar.gz) before training. The manifest fixes the source revision, selected prompt IDs, rejection counts, and output SHA-256 values. Derived dataset files are intentionally not stored in Git. + +Each input row contains one complete preference pair: + +```json +{ + "prompt_id": "stable-id", + "chosen": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}], + "rejected": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}] +} +``` + +Chosen and rejected branches must have an identical prompt and different, non-empty assistant completions. + +## Launch standard DPO + +Download the pinned Qwen checkpoint, then set the model, data, and output locations expected by the standard entrypoint: + +```bash +export MODEL_DIR=/models +export MODEL_REVISION=c1899de289a04d12100db370d81485cdf75e47ca # full 40-character commit SHA +export HF_CHECKPOINT="${MODEL_DIR}/Qwen3-0.6B-${MODEL_REVISION}" +export PROMPT_DATA=/data/task31-ultrafeedback/ultrafeedback_train.parquet +export SAVE_DIR=/checkpoints/task31-dpo + +hf download Qwen/Qwen3-0.6B --revision "${MODEL_REVISION}" --local-dir "${HF_CHECKPOINT}" +bash scripts/training/dpo/run-qwen3-0.6B-ultrafeedback-1xgpu.sh +``` + +The recipe defaults to 200 optimizer steps, 32 preference pairs per global batch, `beta=0.1`, and a 1,024-token branch limit. `GLOBAL_BATCH_SIZE`, `NUM_ROLLOUT`, `MAX_TOKENS_PER_GPU`, and `SAVE_INTERVAL` can be overridden explicitly. + +Standard DPO verifies the pinned repository revision against the local `HF_CHECKPOINT` directory, then reconstructs the frozen reference from that directory. Checkpoints include a reference-identity sidecar containing canonical parameter and fixed-probe digests. A missing or mismatched sidecar fails before the next forward pass. + +The probe digest is a byte-exact SHA-256 over frozen-reference log-probabilities, so resume assumes the same GPU model, driver, image, and kernel stack as the original run. Resuming on different hardware or software fails the probe check by design — treat it as an environment mismatch, not data corruption. + +Use `--dpo-reference-free` only when reference-free DPO is intended; do not combine it with the standard reference identity arguments. + +## Pair-aware batching + +One preference pair is one TransferQueue row. Its chosen and rejected branch lengths are combined into `custom_meta.total_lengths`; the pinned `SeqlenBalancedSampler` assigns complete rows and keeps equal pair counts across data-parallel ranks. Branches are expanded only after a rank receives its rows, so dynamic micro-batch reordering cannot split pair identity. + +## Metrics + +DPO emits the following training metrics under the `train/dpo/` namespace: + +- `loss`, `logps_chosen`, and `logps_rejected`; +- `ref_logps_chosen` and `ref_logps_rejected` in standard mode; +- `reward_chosen`, `reward_rejected`, and `reward_margin`; +- `strict_accuracy`, `tie_rate`, and `tie_aware_accuracy`. + +For distributed parity claims, run DP=1 and DP=2 with the same image, model/data revisions, hyperparameters, and batch semantics, and retain the raw logs and reference digests. diff --git a/docs/zh/guide/dpo-training.md b/docs/zh/guide/dpo-training.md new file mode 100644 index 000000000..068e83630 --- /dev/null +++ b/docs/zh/guide/dpo-training.md @@ -0,0 +1,64 @@ +# DPO 训练 + +Relax 通过离线 SFT 数据链路支持 Direct Preference Optimization（DPO）。Task 31 的公开 recipe 是 [`run-qwen3-0.6B-ultrafeedback-1xgpu.sh`](../../../scripts/training/dpo/run-qwen3-0.6B-ultrafeedback-1xgpu.sh)。 + +## 准备偏好数据子集 + +从固定的数据集 revision 生成确定性的 UltraFeedback 子集： + +```bash +python scripts/data/prepare_ultrafeedback_preferences.py \ + --output-dir /data/task31-ultrafeedback +``` + +命令会生成 train/eval JSONL、Parquet 以及 `manifest.json`。对于已发布的 Task 31 子集，训练前应将生成结果与[可复现性证据包](https://github.com/user-attachments/files/31305744/task31-pr1-dpo-evidence-public-v2.tar.gz)中的 manifest 对比。manifest 固定 source revision、选中的 prompt ID、拒绝原因计数和输出文件 SHA-256；派生数据文件本身不提交到 Git。 + +每行输入承载一个完整 preference pair： + +```json +{ + "prompt_id": "stable-id", + "chosen": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}], + "rejected": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}] +} +``` + +chosen/rejected 必须共享完全相同的 prompt，并包含不同且非空的 assistant completion。 + +## 启动标准 DPO + +下载固定版本的 Qwen checkpoint，然后设置标准入口所需的模型、数据和输出路径： + +```bash +export MODEL_DIR=/models +export MODEL_REVISION=c1899de289a04d12100db370d81485cdf75e47ca # 完整的 40 位 commit SHA +export HF_CHECKPOINT="${MODEL_DIR}/Qwen3-0.6B-${MODEL_REVISION}" +export PROMPT_DATA=/data/task31-ultrafeedback/ultrafeedback_train.parquet +export SAVE_DIR=/checkpoints/task31-dpo + +hf download Qwen/Qwen3-0.6B --revision "${MODEL_REVISION}" --local-dir "${HF_CHECKPOINT}" +bash scripts/training/dpo/run-qwen3-0.6B-ultrafeedback-1xgpu.sh +``` + +recipe 默认运行 200 个 optimizer step，每个 global batch 为 32 个 preference pair，`beta=0.1`，单分支最大 1,024 token。可以显式覆盖 `GLOBAL_BATCH_SIZE`、`NUM_ROLLOUT`、`MAX_TOKENS_PER_GPU` 和 `SAVE_INTERVAL`。 + +标准 DPO 会先在本地 `HF_CHECKPOINT` 目录中校验固定的 repository revision，再从该目录重建冻结 reference。checkpoint 带有 reference identity sidecar，其中保存 canonical parameter digest 和固定 probe digest；sidecar 缺失或不一致时，会在下一次 forward 前失败。 + +probe digest 是对冻结 reference log-probability 的逐字节 SHA-256，因此 resume 假定 GPU 型号、驱动、镜像与内核栈与原运行完全一致。在不同硬件或软件环境上 resume 会按设计触发 probe 校验失败——这表示环境不匹配，而非数据损坏。 + +只有明确需要 reference-free DPO 时才使用 `--dpo-reference-free`，不要同时传入标准 reference identity 参数。 + +## Pair-aware batching + +一个 preference pair 对应一个 TransferQueue row。chosen/rejected 分支长度相加后写入 `custom_meta.total_lengths`；固定版本的 `SeqlenBalancedSampler` 分配完整 row，并保证各 data-parallel rank 的 pair 数相同。只有 rank 收到 pair row 后才展开两个分支，因此动态 micro-batch 重排不会破坏 pair identity。 + +## 指标 + +DPO 在 `train/dpo/` 命名空间下记录以下训练指标： + +- `loss`、`logps_chosen` 和 `logps_rejected`； +- 标准模式下的 `ref_logps_chosen` 和 `ref_logps_rejected`； +- `reward_chosen`、`reward_rejected` 和 `reward_margin`； +- `strict_accuracy`、`tie_rate` 和 `tie_aware_accuracy`。 + +如需声明分布式一致性，应在相同镜像、模型/数据 revision、超参数和 batch 语义下分别运行 DP=1、DP=2，并保留原始日志与 reference digest。 diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 75143efa6..32733cf4d 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -6,6 +6,7 @@ import socket import time from argparse import Namespace +from dataclasses import replace from functools import partial from typing import Any, List @@ -31,6 +32,7 @@ from relax.engine.sft.eval.runner import run_sft_eval from relax.engine.sft.predict.runner import run_sft_predict from relax.engine.sft.runtime import ( + is_preference_mode, is_sft_mode, sft_partition_id, sft_task_name, @@ -77,7 +79,7 @@ from ...utils.profile_utils import TrainProfiler from ...utils.training.tensor_backper import TensorBackuper -from .checkpoint import load_checkpoint +from .checkpoint import is_megatron_checkpoint, load_checkpoint from .collective_utils import _agree_drained from .cp_utils import all_gather_with_cp, maybe_padded_total_lengths, slice_with_cp from .data import ( @@ -85,6 +87,7 @@ DataIterator, build_rollout_minibatch_plan, concat_rollout_batches, + expand_preference_rollout_data, get_data_iterator, log_perf_data, log_perf_data_fwd, @@ -93,6 +96,17 @@ from .initialize import init, is_megatron_main_rank from .loss import compute_advantages_and_returns, get_log_probs_and_entropy, get_values from .model import forward_only, initialize_model_and_optimizer, save, train +from .reference_integrity import ( + REFERENCE_LOADER_MODE, + DPOReferenceIdentity, + canonical_optimizer_sha256, + canonical_tensor_sha256, + read_reference_identity, + reference_identity_path, + reference_probe_sha256, + resolve_dpo_reference_checkpoint, + write_reference_identity, +) from .weight_update.common import named_params_and_buffers from .weight_update.train_offload import MegatronTrainStateOffloader from .weight_update.update_weight_from_distributed import UpdateWeightFromDistributed @@ -226,9 +240,13 @@ def _init( self.args.lr = self.args.critic_lr self.args.lr_warmup_iters = self.args.critic_lr_warmup_iters + resumed_from_megatron = is_megatron_checkpoint(args.load) self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id = initialize_model_and_optimizer( args, role ) + self._dpo_reference_identity: DPOReferenceIdentity | None = None + self._expected_dpo_reference_identity: DPOReferenceIdentity | None = None + self._dpo_reference_probe_verified = False # Train-state offload for colocate sleep/wake. Picks torch_memory_saver # (VMM pause) or manual selective CPU offload based on TMS availability; @@ -262,7 +280,16 @@ def _init( self.weights_backuper.backup("actor") if with_ref: - self.load_other_checkpoint("ref", args.ref_load) + if is_preference_mode(args) and args.sft_objective == "dpo": + reference_checkpoint = resolve_dpo_reference_checkpoint( + args.dpo_reference_repository, args.dpo_reference_revision, args.hf_checkpoint + ) + if resumed_from_megatron: + identity_path = reference_identity_path(args.load, loaded_rollout_id) + self._expected_dpo_reference_identity = read_reference_identity(identity_path) + self._rebuild_dpo_reference(reference_checkpoint) + else: + self.load_other_checkpoint("ref", args.ref_load) # Load teacher model for Megatron-based on-policy distillation if with_opd_teacher: @@ -416,9 +443,186 @@ def _switch_model(self, target_tag: str) -> None: device_utils.maybe_backend_process_on_model_switch() if target_tag not in self.weights_backuper.backup_tags: raise ValueError(f"Cannot switch to unknown model tag: {target_tag}") + if self._active_model_tag == target_tag: + # Same-tag restore would be a byte-identical copy: paths that + # deliberately dirty the weights clear the tag first (see + # _rebuild_dpo_reference), so skipping avoids a redundant + # full-weight CPU->GPU copy per step after the ref forward. + return self.weights_backuper.restore(target_tag) self._active_model_tag = target_tag + def _is_standard_dpo(self) -> bool: + return is_preference_mode(self.args) and self.args.sft_objective == "dpo" and not self.args.dpo_reference_free + + def _assert_dp_reference_digest_equal(self, digest: str) -> None: + digests = [None] * dist.get_world_size(group=get_gloo_group()) + dist.all_gather_object(digests, digest, group=get_gloo_group()) + if len(set(digests)) != 1: + raise RuntimeError(f"DPO frozen-reference parameter digests differ across ranks: {digests}") + + def _assert_dpo_reference_identity(self, actual: DPOReferenceIdentity) -> None: + expected = self._expected_dpo_reference_identity + if expected is None: + return + fields = ("repository", "revision", "loader_mode", "parameter_sha256") + mismatches = { + field: (getattr(expected, field), getattr(actual, field)) + for field in fields + if getattr(expected, field) != getattr(actual, field) + } + if mismatches: + raise RuntimeError(f"DPO frozen-reference identity mismatch: {mismatches}") + + def _rebuild_dpo_reference(self, path: str) -> None: + """Transactionally rebuild a frozen reference without touching + optimizer state.""" + if self._active_model_tag != "actor" or "actor" not in self.weights_backuper.backup_tags: + raise RuntimeError("DPO reference rebuild requires an active actor backup") + optimizer_before = canonical_optimizer_sha256(self.optimizer) + old_args = self.args.load, self.args.no_load_optim, self.args.no_load_rng, self.args.finetune + try: + self.args.load = path + self.args.no_load_optim = True + self.args.no_load_rng = True + self.args.finetune = True + self._active_model_tag = None + load_checkpoint( + self.model, + None, + None, + checkpointing_context={}, + skip_load_to_model_and_opt=False, + ) + candidate_sha256 = canonical_tensor_sha256( + named_params_and_buffers( + self.args, + self.model, + convert_to_global_name=self.args.megatron_to_hf_mode == "raw", + translate_gpu_to_cpu=True, + ) + ) + self._assert_dp_reference_digest_equal(candidate_sha256) + candidate = DPOReferenceIdentity( + schema_version=1, + repository=self.args.dpo_reference_repository, + revision=self.args.dpo_reference_revision, + loader_mode=REFERENCE_LOADER_MODE, + parameter_sha256=candidate_sha256, + probe_sha256=None, + probe_manifest=None, + ) + self._assert_dpo_reference_identity(candidate) + self.weights_backuper.backup("ref") + self._dpo_reference_identity = candidate + finally: + self.args.load, self.args.no_load_optim, self.args.no_load_rng, self.args.finetune = old_args + self._switch_model("actor") + optimizer_after = canonical_optimizer_sha256(self.optimizer) + if optimizer_after != optimizer_before: + raise RuntimeError( + "DPO reference rebuild modified optimizer master parameters or state: " + f"before={optimizer_before}, after={optimizer_after}" + ) + + def _validate_dpo_reference_probe(self, rollout_data: RolloutBatch) -> None: + if self._expected_dpo_reference_identity is not None: + if not self._dpo_reference_probe_verified: + raise RuntimeError("resumed DPO reference probe must be replayed before training data forward") + return + if self._dpo_reference_identity is not None and self._dpo_reference_identity.probe_sha256 is not None: + return + first_pair_id = int(rollout_data["preference_branch_pair_ids"][0]) + indices = [ + index + for index, pair_id in enumerate(rollout_data["preference_branch_pair_ids"]) + if int(pair_id) == first_pair_id + ] + if len(indices) != 2: + raise RuntimeError(f"DPO reference probe pair {first_pair_id!r} is not atomic") + manifest = { + "pair_ids": [int(rollout_data["preference_branch_pair_ids"][index]) for index in indices], + "branch_is_chosen": [bool(rollout_data["preference_is_chosen"][index]) for index in indices], + "tokens": [torch.as_tensor(rollout_data["tokens"][index]).cpu().tolist() for index in indices], + "loss_masks": [torch.as_tensor(rollout_data["loss_masks"][index]).cpu().tolist() for index in indices], + "total_lengths": [int(rollout_data["total_lengths"][index]) for index in indices], + "response_lengths": [int(rollout_data["response_lengths"][index]) for index in indices], + } + if self._dpo_reference_identity is None: + raise RuntimeError("DPO frozen-reference identity was not initialized") + probe_sha256 = self._compute_dpo_reference_probe(manifest) + self._dpo_reference_identity = replace( + self._dpo_reference_identity, probe_sha256=probe_sha256, probe_manifest=manifest + ) + + def _compute_dpo_reference_probe(self, manifest: dict[str, Any]) -> str: + """Run the canonical single-pair probe without perturbing training + RNG.""" + pair_ids = [int(value) for value in manifest["pair_ids"]] + branch_is_chosen = [bool(value) for value in manifest["branch_is_chosen"]] + if len(pair_ids) != 2 or set(branch_is_chosen) != {False, True} or len(set(pair_ids)) != 1: + raise RuntimeError("DPO reference probe manifest must contain one atomic preference pair") + device = device_utils.make_current_torch_device() + tokens = [torch.tensor(value, dtype=torch.long, device=device) for value in manifest["tokens"]] + loss_masks = [torch.tensor(value, dtype=torch.bool, device=device) for value in manifest["loss_masks"]] + total_lengths = [int(value) for value in manifest["total_lengths"]] + response_lengths = [int(value) for value in manifest["response_lengths"]] + dp_size = mpu.get_data_parallel_world_size(with_context_parallel=False) + probe_data: RolloutBatch = { + "tokens": tokens, + "loss_masks": loss_masks, + "total_lengths": total_lengths, + "response_lengths": response_lengths, + "preference_branch_pair_ids": pair_ids, + "preference_is_chosen": branch_is_chosen, + "preference_pair_ids": [pair_ids[0]], + "preference_pair_costs": [sum(total_lengths)], + "dynamic_global_batch_size": dp_size, + } + probe_iterator, probe_microbatches = get_data_iterator(self.args, self.model, probe_data) + restore_tag = self._active_model_tag + if restore_tag is None: + raise RuntimeError("DPO reference probe requires an active model backup") + python_rng_state = random.getstate() + torch_rng_state = torch.get_rng_state() + cuda_rng_states = torch.cuda.get_rng_state_all() + try: + self._switch_model("ref") + output = self.compute_log_prob(probe_iterator, probe_microbatches, store_prefix="ref_") + finally: + self._switch_model(restore_tag) + random.setstate(python_rng_state) + torch.set_rng_state(torch_rng_state) + torch.cuda.set_rng_state_all(cuda_rng_states) + return reference_probe_sha256( + pair_ids, + branch_is_chosen, + tokens, + loss_masks, + output["ref_log_probs"], + ) + + def _replay_dpo_reference_probe(self) -> None: + expected = self._expected_dpo_reference_identity + if expected is None or self._dpo_reference_probe_verified: + return + manifest = expected.probe_manifest + if expected.probe_sha256 is None or manifest is None: + raise RuntimeError("resumed DPO checkpoint is missing frozen-reference probe metadata") + actual = self._compute_dpo_reference_probe(manifest) + if actual != expected.probe_sha256: + raise RuntimeError( + f"DPO frozen-reference probe mismatch: expected={expected.probe_sha256}, actual={actual}" + ) + if self._dpo_reference_identity is None: + raise RuntimeError("DPO frozen-reference identity was not initialized") + self._dpo_reference_identity = replace( + self._dpo_reference_identity, + probe_sha256=expected.probe_sha256, + probe_manifest=manifest, + ) + self._dpo_reference_probe_verified = True + def fill_routing_replay(self, data_iterator, num_microbatches, rollout_data): if "rollout_routed_experts" not in rollout_data: raise ValueError( @@ -783,6 +987,9 @@ def train_critic(self, rollout_id: int, rollout_data: RolloutBatch) -> None: self.sleep() def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: + if is_preference_mode(self.args): + rollout_data = expand_preference_rollout_data(rollout_data) + # PPO colocate: ``values`` and ``loss_masks`` reach us via TransferQueue # and land on CPU (critic ``.cpu()`` s ``values`` before PUT). Inline # GAE + normalize_advantages need GPU tensors — dispatch here so the @@ -799,6 +1006,8 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: # Create data iterator for actor forward + routing replay + train. data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data) + if self._is_standard_dpo(): + self._replay_dpo_reference_probe() # Create a separate iterator with a larger token budget for ref/teacher log-probs if self.args.use_dynamic_batch_size and self.args.log_probs_max_tokens_per_gpu != self.args.max_tokens_per_gpu: data_iterator_logprobs, num_microbatches_logprobs = get_data_iterator( @@ -815,7 +1024,10 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: with inverse_timer("train_wait"), timer("train"): # All RL algorithms need ref/teacher/actor inline forwards to produce old_log_probs. - should_compute_old_log_probs = self.args.compute_advantages_and_returns + standard_dpo = ( + is_preference_mode(self.args) and self.args.sft_objective == "dpo" and not self.args.dpo_reference_free + ) + should_compute_old_log_probs = self.args.compute_advantages_and_returns or standard_dpo # PPO fully_async has a standalone Advantages service that produces # advantages/returns via TransferQueue; every other path (including # PPO colocate) computes GAE inline from critic's ``values``. @@ -827,14 +1039,19 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: if "ref" in self.weights_backuper.backup_tags: if self.args.use_routing_replay: os.environ["ROUTING_REPLAY_STAGE"] = "fallthrough" - self._switch_model("ref") - rollout_data.update( - self.compute_log_prob( - data_iterator_logprobs, - num_microbatches_logprobs, - store_prefix="ref_", + try: + self._switch_model("ref") + rollout_data.update( + self.compute_log_prob( + data_iterator_logprobs, + num_microbatches_logprobs, + store_prefix="ref_", + ) ) - ) + finally: + self._switch_model("old_actor" if self.args.keep_old_actor else "actor") + if standard_dpo: + self._validate_dpo_reference_probe(rollout_data) # Forward teacher model to get teacher_log_probs for Megatron-based OPD if "teacher" in self.weights_backuper.backup_tags: @@ -851,7 +1068,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: ) self._switch_model("old_actor" if self.args.keep_old_actor else "actor") - if not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics: + if not standard_dpo and (not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics): if self.args.use_routing_replay: if self.args.use_rollout_routing_replay: os.environ["ROUTING_REPLAY_STAGE"] = "replay_forward" @@ -1596,9 +1813,26 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: save(rollout_id, self.model, self.optimizer, self.opt_param_scheduler) - if force_sync and self.args.async_save: + if (force_sync or self._is_standard_dpo()) and self.args.async_save: maybe_finalize_async_save(blocking=True) + if self._is_standard_dpo(): + identity = self._dpo_reference_identity + if identity is None or identity.probe_sha256 is None: + raise RuntimeError("cannot save standard DPO without a validated reference identity and probe") + actual_sha256 = canonical_tensor_sha256(self.weights_backuper.get("ref").items()) + if actual_sha256 != identity.parameter_sha256: + raise RuntimeError( + "DPO frozen-reference checksum changed before checkpoint: " + f"expected={identity.parameter_sha256}, actual={actual_sha256}" + ) + if dist.get_rank(group=get_gloo_group()) == 0: + write_reference_identity(reference_identity_path(self.args.save, rollout_id), identity) + # Functional barrier (not debugging): peers must not proceed past + # the checkpoint before rank 0's identity sidecar is durable, + # otherwise a concurrent resume could miss the file. + dist.barrier(group=get_gloo_group()) + if self.args.save_hf is not None and self.role == "actor": from relax.backends.megatron.model import save_hf_model @@ -1852,30 +2086,31 @@ def recv_weight_fully_async(self, rollout_id) -> None: def load_other_checkpoint(self, model_tag: str, path: str) -> None: old_args = self.args.load, self.args.no_load_optim, self.args.no_load_rng, self.args.finetune - self.args.load = path - self.args.no_load_optim = True - self.args.no_load_rng = True - self.args.finetune = True - old_ckpt_step = None - if model_tag == "ref" and self.args.ref_ckpt_step is not None: - old_ckpt_step = self.args.ckpt_step - self.args.ckpt_step = self.args.ref_ckpt_step - elif model_tag == "teacher" and self.args.opd_teacher_ckpt_step is not None: - old_ckpt_step = self.args.ckpt_step - self.args.ckpt_step = self.args.opd_teacher_ckpt_step - - _, _ = load_checkpoint( - self.model, - None, - None, - checkpointing_context={}, - skip_load_to_model_and_opt=False, - ) - self.args.load, self.args.no_load_optim, self.args.no_load_rng, self.args.finetune = old_args - - if old_ckpt_step is not None: - self.args.ckpt_step = old_ckpt_step + try: + self.args.load = path + self.args.no_load_optim = True + self.args.no_load_rng = True + self.args.finetune = True + + if model_tag == "ref" and self.args.ref_ckpt_step is not None: + old_ckpt_step = self.args.ckpt_step + self.args.ckpt_step = self.args.ref_ckpt_step + elif model_tag == "teacher" and self.args.opd_teacher_ckpt_step is not None: + old_ckpt_step = self.args.ckpt_step + self.args.ckpt_step = self.args.opd_teacher_ckpt_step + + _, _ = load_checkpoint( + self.model, + None, + None, + checkpointing_context={}, + skip_load_to_model_and_opt=False, + ) + finally: + self.args.load, self.args.no_load_optim, self.args.no_load_rng, self.args.finetune = old_args + if old_ckpt_step is not None: + self.args.ckpt_step = old_ckpt_step self.weights_backuper.backup(model_tag) self._active_model_tag = model_tag diff --git a/relax/backends/megatron/checkpoint.py b/relax/backends/megatron/checkpoint.py index 76752d2f1..3ccb09699 100644 --- a/relax/backends/megatron/checkpoint.py +++ b/relax/backends/megatron/checkpoint.py @@ -130,7 +130,7 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, checkpointing_con exist = Path(load_path).exists() and _is_dir_nonempty(load_path) - if exist and _is_megatron_checkpoint(load_path): + if exist and is_megatron_checkpoint(load_path): _alias_renamed_transfer_queue_enum() try: return _load_checkpoint_megatron( @@ -190,7 +190,9 @@ def _format_opt_param_scheduler_error(args, original: AssertionError) -> str: ) -def _is_megatron_checkpoint(path: str | Path) -> bool: +def is_megatron_checkpoint(path: str | Path | None) -> bool: + if path is None: + return False return (Path(path) / "latest_checkpointed_iteration.txt").is_file() or bool( re.fullmatch(r"iter_\d{7}", Path(path).name) ) diff --git a/relax/backends/megatron/data.py b/relax/backends/megatron/data.py index 9aff3da6f..66ef35afb 100644 --- a/relax/backends/megatron/data.py +++ b/relax/backends/megatron/data.py @@ -15,6 +15,7 @@ from megatron.training.global_vars import get_args from torch.nn.utils.rnn import pad_sequence +from relax.engine.sft.runtime import is_preference_mode from relax.utils import device as device_utils from relax.utils import tracking_utils from relax.utils.data.data import get_minimum_num_micro_batch_size @@ -25,6 +26,7 @@ from relax.utils.timer import Timer from relax.utils.training import train_metric_utils from relax.utils.training.flops_counter import FlopsCounter +from relax.utils.training.preference_utils import pack_preference_pair_indices from relax.utils.types import RolloutBatch from .cp_utils import ( @@ -703,6 +705,139 @@ def reset(self) -> "DataIterator": return self +def expand_preference_rollout_data(rollout_data: RolloutBatch) -> RolloutBatch: + """Expand pair rows into adjacent chosen/rejected model sequences.""" + if "preference_pair_costs" in rollout_data: + return rollout_data + required = ( + "pair_ids", + "chosen_tokens", + "rejected_tokens", + "chosen_loss_masks", + "rejected_loss_masks", + "chosen_total_lengths", + "rejected_total_lengths", + ) + missing = [key for key in required if key not in rollout_data] + if missing: + raise ValueError(f"preference rollout data is missing fields: {missing}") + pair_count = len(rollout_data["pair_ids"]) + if pair_count <= 0: + raise ValueError("preference rollout batch must contain at least one pair") + for key in required: + if len(rollout_data[key]) != pair_count: + raise ValueError( + f"preference field {key!r} is not pair-row aligned: expected {pair_count}, got {len(rollout_data[key])}" + ) + + flat: RolloutBatch = { + "tokens": [], + "loss_masks": [], + "total_lengths": [], + "response_lengths": [], + "preference_branch_pair_ids": [], + "preference_is_chosen": [], + } + pair_costs: list[int] = [] + pair_ids: list[int] = [] + for index in range(pair_count): + chosen_length = int(rollout_data["chosen_total_lengths"][index]) + rejected_length = int(rollout_data["rejected_total_lengths"][index]) + if chosen_length <= 0 or rejected_length <= 0: + raise ValueError(f"preference pair row {index} has non-positive branch length") + pair_id = int(rollout_data["pair_ids"][index]) + pair_ids.append(pair_id) + pair_costs.append(chosen_length + rejected_length) + for prefix, is_chosen in (("chosen", True), ("rejected", False)): + tokens = rollout_data[f"{prefix}_tokens"][index] + loss_mask = rollout_data[f"{prefix}_loss_masks"][index] + total_length = int(rollout_data[f"{prefix}_total_lengths"][index]) + if len(tokens) != total_length or len(loss_mask) != total_length: + raise ValueError( + f"preference pair row {index} {prefix} tensor length does not match declared total_length" + ) + flat["tokens"].append(tokens) + flat["loss_masks"].append(loss_mask) + flat["total_lengths"].append(total_length) + flat["response_lengths"].append(total_length) + flat["preference_branch_pair_ids"].append(pair_id) + flat["preference_is_chosen"].append(is_chosen) + flat["preference_pair_costs"] = pair_costs + flat["preference_pair_ids"] = pair_ids + if "dynamic_global_batch_size" in rollout_data: + flat["dynamic_global_batch_size"] = rollout_data["dynamic_global_batch_size"] + return flat + + +def _split_preference_bins_to_count(bins: list[list[int]], target_count: int) -> list[list[int]]: + bins = [list(group) for group in bins] + while len(bins) < target_count: + candidates = [(len(group), -index, index) for index, group in enumerate(bins) if len(group) > 1] + if not candidates: + raise RuntimeError( + f"cannot split {len(bins)} preference micro-batches to DP-synchronized count {target_count}" + ) + _, _, index = max(candidates) + group = bins[index] + bins[index] = group[:-1] + bins.insert(index + 1, [group[-1]]) + return bins + + +def _get_preference_data_iterator( + args: Namespace, + rollout_data: RolloutBatch, + max_tokens_per_gpu: int | None, +) -> tuple[list[DataIterator], list[int]]: + pair_costs = [int(cost) for cost in rollout_data["preference_pair_costs"]] + pair_ids = [str(pair_id) for pair_id in rollout_data["preference_pair_ids"]] + dp_size = mpu.get_data_parallel_world_size(with_context_parallel=False) + dynamic_count = rollout_data.get("dynamic_global_batch_size") + if isinstance(dynamic_count, (list, tuple)): + normalized = {int(value) for value in dynamic_count} + if len(normalized) != 1: + raise ValueError(f"dynamic_global_batch_size values must be identical, got {dynamic_count}") + dynamic_count = normalized.pop() + expected_global_pairs = int(args.global_batch_size) if dynamic_count is None else int(dynamic_count) + capacity = int(max_tokens_per_gpu or args.max_tokens_per_gpu) + pair_bins = pack_preference_pair_indices(pair_costs, pair_ids, capacity=capacity) + control_values = [None] * dp_size + dist.all_gather_object( + control_values, + (len(pair_costs), expected_global_pairs, len(pair_bins)), + group=mpu.get_data_parallel_group_gloo(with_context_parallel=False), + ) + local_pair_counts = {int(values[0]) for values in control_values} + if len(local_pair_counts) != 1: + raise ValueError(f"preference objectives require equal local pair rows on every DP rank: {local_pair_counts}") + declared_global_pairs = {int(values[1]) for values in control_values} + if len(declared_global_pairs) != 1: + raise ValueError(f"dynamic_global_batch_size must be identical on every DP rank: {declared_global_pairs}") + step_global_pair_count = local_pair_counts.pop() * dp_size + declared_global_pair_count = declared_global_pairs.pop() + if declared_global_pair_count != step_global_pair_count: + raise ValueError( + "dynamic_global_batch_size must equal the step-global preference pair count: " + f"declared={declared_global_pair_count}, actual={step_global_pair_count}, " + f"local={len(pair_costs)}, dp_size={dp_size}" + ) + rollout_data["dynamic_global_batch_size"] = step_global_pair_count + pair_bins = _split_preference_bins_to_count(pair_bins, max(int(values[2]) for values in control_values)) + if any(sum(pair_costs[index] for index in group) > capacity for group in pair_bins): + raise RuntimeError("preference DP bin synchronization produced an over-capacity micro-batch") + branch_bins = [ + [branch for pair_index in group for branch in (2 * pair_index, 2 * pair_index + 1)] for group in pair_bins + ] + covered = [index for group in pair_bins for index in group] + if sorted(covered) != list(range(len(pair_costs))): + raise RuntimeError("preference dynamic batching lost or duplicated a pair") + for group in branch_bins: + if len(group) % 2 != 0 or any(group[index + 1] != group[index] + 1 for index in range(0, len(group), 2)): + raise RuntimeError("preference dynamic batching split a chosen/rejected pair") + iterator = DataIterator(rollout_data, micro_batch_indices=branch_bins, max_tokens_per_gpu=capacity) + return [iterator], [len(branch_bins)] + + def get_data_iterator( args: Namespace, model: torch.nn.Module | Sequence[torch.nn.Module], @@ -722,6 +857,9 @@ def get_data_iterator( - `data_iterators`: list of `DataIterator`, one per VPP stage (size 1 if VPP disabled) - `num_microbatches`: list[int], one per local step in the rollout (length = steps) """ + if is_preference_mode(args): + return _get_preference_data_iterator(args, rollout_data, max_tokens_per_gpu) + dp_size = mpu.get_data_parallel_world_size(with_context_parallel=False) dp_group = mpu.get_data_parallel_group() vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() @@ -917,6 +1055,10 @@ def log_rollout_data( ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY, ROLLOUT_MINI_GLOBAL_SAMPLE_COUNTS_KEY, ROLLOUT_MINI_PROMPT_GROUP_COUNTS_KEY, + "preference_pair_costs", + "preference_pair_ids", + "preference_branch_pair_ids", + "preference_is_chosen", ]: continue if args.use_opd and key in OPD_ROLLOUT_LOG_SKIP_FIELDS: diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py index 170265a24..6bcf4a26f 100644 --- a/relax/backends/megatron/loss.py +++ b/relax/backends/megatron/loss.py @@ -9,6 +9,7 @@ from megatron.core import mpu from torch.utils.checkpoint import checkpoint +from relax.engine.sft.runtime import is_preference_mode from relax.utils.distributed_utils import distributed_masked_normalize, distributed_masked_whiten from relax.utils.misc import load_function from relax.utils.opd.opd_utils import ( @@ -33,6 +34,11 @@ get_reinforce_plus_plus_baseline_advantages, get_reinforce_plus_plus_returns, ) +from relax.utils.training.preference_utils import ( + build_preference_pair_indices, + dpo_pair_loss, + require_tensor_condition, +) from relax.utils.types import RolloutBatch from .cp_utils import ( @@ -1265,6 +1271,108 @@ def sft_loss_function( ) +def dpo_loss_function( + args: Namespace, + batch: RolloutBatch, + logits: torch.Tensor, + sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor], # noqa: ARG001 +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Compute a pair-summed DPO objective using explicit pair identity.""" + if len(batch["response_lengths"]) % 2 != 0: + raise ValueError("DPO micro-batch must contain an even number of chosen/rejected branches") + _, values = get_log_probs_and_entropy( + logits, + args=args, + unconcat_tokens=batch["unconcat_tokens"], + total_lengths=batch["total_lengths"], + response_lengths=batch["response_lengths"], + with_entropy=False, + max_seq_lens=batch.get("max_seq_lens", None), + padded_total_lengths=batch.get("padded_total_lengths", None), + dynamic_cp_size=batch.get("dynamic_cp_size", None), + dynamic_cp_rank=batch.get("dynamic_cp_rank", None), + ) + policy_token_log_probs = values["log_probs"] + + def sequence_sums(token_values) -> torch.Tensor: + if len(token_values) != len(batch["loss_masks"]): + raise ValueError("DPO token log-probabilities are not branch aligned") + sums = [] + for branch_values, mask in zip(token_values, batch["loss_masks"], strict=True): + branch_values = torch.as_tensor(branch_values, device=logits.device) + branch_mask = mask.to(device=logits.device, dtype=branch_values.dtype) + if branch_values.shape != branch_mask.shape: + raise ValueError( + "DPO branch log-probability/mask shape mismatch: " + f"{tuple(branch_values.shape)} vs {tuple(branch_mask.shape)}" + ) + require_tensor_condition( + branch_mask.to(dtype=torch.bool).any(), + "DPO branch completion mask must contain at least one supervised token", + ) + sums.append((branch_values * branch_mask).sum()) + return torch.stack(sums) + + pair_ids = batch.get("preference_branch_pair_ids") + branch_is_chosen = batch.get("preference_is_chosen") + if pair_ids is None or branch_is_chosen is None: + raise ValueError("DPO batch is missing preference pair identity fields") + chosen_indices, rejected_indices = build_preference_pair_indices(pair_ids, branch_is_chosen) + chosen_index = torch.as_tensor(chosen_indices, dtype=torch.long, device=logits.device) + rejected_index = torch.as_tensor(rejected_indices, dtype=torch.long, device=logits.device) + + policy_sums = sequence_sums(policy_token_log_probs) + policy_chosen = policy_sums.index_select(0, chosen_index) + policy_rejected = policy_sums.index_select(0, rejected_index) + reference_free = bool(args.dpo_reference_free) + if reference_free: + reference_chosen = reference_rejected = None + ref_chosen_for_metrics = torch.zeros_like(policy_chosen) + ref_rejected_for_metrics = torch.zeros_like(policy_rejected) + else: + reference_values = batch.get("ref_log_probs") + if reference_values is None: + raise ValueError("standard DPO batch is missing frozen-reference log-probabilities") + reference_sums = sequence_sums(reference_values) + reference_chosen = reference_sums.index_select(0, chosen_index) + reference_rejected = reference_sums.index_select(0, rejected_index) + ref_chosen_for_metrics = reference_chosen + ref_rejected_for_metrics = reference_rejected + pair_losses = dpo_pair_loss( + policy_chosen, + policy_rejected, + reference_chosen=reference_chosen, + reference_rejected=reference_rejected, + beta=args.dpo_beta, + reference_free=reference_free, + ) + chosen_rewards = args.dpo_beta * (policy_chosen - ref_chosen_for_metrics) + rejected_rewards = args.dpo_beta * (policy_rejected - ref_rejected_for_metrics) + # pair_losses is never empty: build_preference_pair_indices raises on an + # empty micro-batch, so no gradient-safety fallback is needed here. + loss = pair_losses.sum() + reward_margin = chosen_rewards - rejected_rewards + tie = reward_margin.abs() <= 1e-6 + strict = reward_margin > 0 + correct = reward_margin > 1e-6 + metrics = { + "dpo/loss": pair_losses.detach().sum(), + "dpo/logps_chosen": policy_chosen.detach().sum(), + "dpo/logps_rejected": policy_rejected.detach().sum(), + "dpo/reward_chosen": chosen_rewards.detach().sum(), + "dpo/reward_rejected": rejected_rewards.detach().sum(), + "dpo/reward_margin": reward_margin.detach().sum(), + "dpo/strict_accuracy": strict.to(torch.float32).detach().sum(), + "dpo/tie_rate": tie.to(torch.float32).detach().sum(), + "dpo/tie_aware_accuracy": (correct.to(torch.float32) + 0.5 * tie.to(torch.float32)).detach().sum(), + } + metrics["dpo/pair_accuracy"] = metrics["dpo/strict_accuracy"] + if not reference_free: + metrics["dpo/ref_logps_chosen"] = reference_chosen.detach().sum() + metrics["dpo/ref_logps_rejected"] = reference_rejected.detach().sum() + return loss, metrics + + def sft_loss_function_chunked( args: Namespace, batch: RolloutBatch, @@ -1364,6 +1472,10 @@ def loss_function( dynamic_cp_rank=batch.get("dynamic_cp_rank", None), ) num_samples = len(batch["response_lengths"]) + if is_preference_mode(args): + if num_samples % 2 != 0: + raise ValueError("preference micro-batch contains an odd number of branches") + num_samples //= 2 sum_of_sample_mean = get_sum_of_sample_mean( batch["total_lengths"], @@ -1383,7 +1495,9 @@ def loss_function( case "value_loss": func = value_loss_function case "sft": - if getattr(args, "sft_chunked_logits", False) and lm_head_forward is not None: + if getattr(args, "sft_objective", "causal_lm") == "dpo": + func = dpo_loss_function + elif getattr(args, "sft_chunked_logits", False) and lm_head_forward is not None: # Bind lm_head_forward so chunked path matches the standard # inner-func signature; outer body (recompute, CP guard, # Megatron scaling, return-tuple) is then shared with legacy. diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 077e461d9..36960f738 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -1024,6 +1024,8 @@ def forward_step( "loss_masks", "log_probs", "ref_log_probs", + "preference_branch_pair_ids", + "preference_is_chosen", "values", "advantages", "returns", diff --git a/relax/backends/megatron/reference_integrity.py b/relax/backends/megatron/reference_integrity.py new file mode 100644 index 000000000..f9b989738 --- /dev/null +++ b/relax/backends/megatron/reference_integrity.py @@ -0,0 +1,306 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""DPO frozen-reference identity and byte-level integrity helpers.""" + +import hashlib +import json +import math +import os +import re +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any + +import torch + + +REFERENCE_IDENTITY_FILENAME = "relax_dpo_reference.json" +REFERENCE_LOADER_MODE = "hf_bridge_model_only_v1" +_GIT_COMMIT_SHA256_RE = re.compile(r"^[0-9a-f]{40}$") +_HF_ETAG_RE = re.compile(r"^(?:[0-9a-f]{40}|[0-9a-f]{64})$") +_WEIGHT_INDEXES = ( + ("model.safetensors.index.json", ".safetensors"), + ("pytorch_model.bin.index.json", ".bin"), +) +_SINGLE_WEIGHT_NAMES = ("model.safetensors", "pytorch_model.bin") + + +def _file_matches_hf_etag(path: Path, etag: str) -> bool: + digest = hashlib.sha256() if len(etag) == 64 else hashlib.sha1() + with path.open("rb") as file: + if len(etag) == 40: + size = os.fstat(file.fileno()).st_size + digest.update(f"blob {size}\0".encode()) + while chunk := file.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() == etag + + +def _validate_reference_weight_files(checkpoint: Path) -> None: + found_weights = any((checkpoint / filename).is_file() for filename in _SINGLE_WEIGHT_NAMES) + for index_name, shard_suffix in _WEIGHT_INDEXES: + index_path = checkpoint / index_name + if not index_path.is_file(): + continue + found_weights = True + try: + payload = json.loads(index_path.read_text(encoding="utf-8")) + weight_map = payload["weight_map"] + except (OSError, ValueError, KeyError, TypeError) as exc: + raise RuntimeError(f"DPO reference weight index is invalid: {index_name!r}") from exc + if not isinstance(weight_map, Mapping) or not weight_map: + raise RuntimeError(f"DPO reference weight index has an empty weight_map: {index_name!r}") + shard_names = set() + for filename in weight_map.values(): + if not isinstance(filename, str) or not filename: + raise RuntimeError(f"DPO reference weight index contains an invalid shard name: {index_name!r}") + shard_names.add(filename) + for filename in shard_names: + relative_path = Path(filename) + if relative_path.is_absolute() or ".." in relative_path.parts: + raise RuntimeError(f"DPO reference weight index contains an unsafe shard path: {filename!r}") + if relative_path.suffix != shard_suffix: + raise RuntimeError(f"DPO reference weight index contains an invalid shard suffix: {filename!r}") + if not (checkpoint / relative_path).is_file(): + raise RuntimeError(f"DPO reference weight index points to a missing shard: {filename!r}") + if not found_weights: + supported = ", ".join((*_SINGLE_WEIGHT_NAMES, *(name for name, _suffix in _WEIGHT_INDEXES))) + raise RuntimeError(f"DPO reference has no supported model weights or index; expected one of: {supported}") + + +def _validate_local_download_metadata(checkpoint: Path, revision: str) -> None: + metadata_root = checkpoint / ".cache" / "huggingface" / "download" + snapshot_files = [ + path for path in checkpoint.rglob("*") if path.is_file() and checkpoint / ".cache" not in path.parents + ] + for path in snapshot_files: + relative_path = path.relative_to(checkpoint) + metadata_path = metadata_root.joinpath(*relative_path.parts).with_name(f"{relative_path.name}.metadata") + try: + lines = metadata_path.read_text(encoding="utf-8").splitlines() + metadata_revision, etag, timestamp_text = lines + timestamp = float(timestamp_text) + except (OSError, ValueError) as exc: + raise RuntimeError( + f"DPO reference file is missing valid Hugging Face local-dir metadata: {relative_path.as_posix()!r}" + ) from exc + if ( + metadata_revision != revision + or _HF_ETAG_RE.fullmatch(etag) is None + or not math.isfinite(timestamp) + or path.stat().st_mtime - 1 > timestamp + ): + raise RuntimeError( + "DPO reference file metadata does not match the pinned revision or file contents: " + f"{relative_path.as_posix()!r}" + ) + if not _file_matches_hf_etag(path, etag): + raise RuntimeError( + f"DPO reference file contents do not match its Hugging Face ETag: {relative_path.as_posix()!r}" + ) + + +def resolve_dpo_reference_checkpoint(repository: str, revision: str, local_checkpoint: str) -> str: + """Resolve a pinned DPO reference in the configured local model + directory.""" + try: + from huggingface_hub import snapshot_download + except ImportError as exc: + raise RuntimeError("standard DPO requires huggingface_hub to resolve its frozen reference") from exc + + if _GIT_COMMIT_SHA256_RE.fullmatch(revision) is None: + raise ValueError("standard DPO requires --dpo-reference-revision to be a full 40-character commit SHA") + + try: + configured_checkpoint = Path(local_checkpoint).resolve(strict=True) + checkpoint = Path( + snapshot_download( + repo_id=repository, + revision=revision, + local_dir=str(configured_checkpoint), + local_files_only=True, + ) + ).resolve(strict=True) + except OSError as exc: + raise RuntimeError( + "standard DPO requires its pinned reference in --hf-checkpoint; prepare it with " + f"`hf download {repository} --revision {revision} --local-dir {local_checkpoint}`" + ) from exc + if checkpoint != configured_checkpoint: + raise RuntimeError( + "DPO reference resolution returned a directory different from --hf-checkpoint: " + f"configured={configured_checkpoint}, resolved={checkpoint}" + ) + if not (checkpoint / "config.json").is_file(): + raise RuntimeError( + "resolved DPO reference snapshot is missing config.json: " + f"repository={repository!r}, revision={revision!r}, path={checkpoint}" + ) + try: + _validate_reference_weight_files(checkpoint) + _validate_local_download_metadata(checkpoint, revision) + except RuntimeError as exc: + raise RuntimeError( + f"{exc}; re-download with `hf download {repository} --revision {revision} --local-dir {local_checkpoint}`" + ) from exc + return str(checkpoint) + + +@dataclass(frozen=True) +class DPOReferenceIdentity: + """Identity persisted beside every standard-DPO checkpoint.""" + + schema_version: int + repository: str + revision: str + loader_mode: str + parameter_sha256: str + probe_sha256: str | None + probe_manifest: dict[str, Any] | None = None + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> "DPOReferenceIdentity": + return cls( + schema_version=int(value["schema_version"]), + repository=str(value["repository"]), + revision=str(value["revision"]), + loader_mode=str(value["loader_mode"]), + parameter_sha256=str(value["parameter_sha256"]), + probe_sha256=None if value.get("probe_sha256") is None else str(value["probe_sha256"]), + probe_manifest=None if value.get("probe_manifest") is None else dict(value["probe_manifest"]), + ) + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +def _update_field(digest: Any, value: bytes) -> None: + digest.update(len(value).to_bytes(8, "big")) + digest.update(value) + + +def _tensor_bytes(tensor: torch.Tensor) -> bytes: + value = tensor.detach().cpu().contiguous().reshape(-1) + return value.view(torch.uint8).numpy().tobytes() + + +def canonical_tensor_sha256(named_tensors: Iterable[tuple[str, torch.Tensor]]) -> str: + """Hash names, dtype, shape and bytes in canonical name order.""" + normalized = sorted(((str(name), tensor) for name, tensor in named_tensors), key=lambda item: item[0]) + if not normalized: + raise ValueError("canonical tensor digest requires at least one tensor") + digest = hashlib.sha256() + for name, tensor in normalized: + _update_field(digest, name.encode()) + _update_field(digest, str(tensor.dtype).encode()) + _update_field(digest, json.dumps(list(tensor.shape), separators=(",", ":")).encode()) + _update_field(digest, _tensor_bytes(tensor)) + return digest.hexdigest() + + +def canonical_optimizer_sha256(optimizer: Any) -> str: + """Hash optimizer master parameters and state without relying on object + IDs.""" + digest = hashlib.sha256() + + def update(value: Any, path: str) -> None: + _update_field(digest, path.encode()) + if isinstance(value, torch.Tensor): + _update_field(digest, str(value.dtype).encode()) + _update_field(digest, json.dumps(list(value.shape), separators=(",", ":")).encode()) + _update_field(digest, _tensor_bytes(value)) + elif isinstance(value, Mapping): + for key in sorted(value, key=lambda item: str(item)): + update(value[key], f"{path}/{key}") + elif isinstance(value, (list, tuple)): + for index, item in enumerate(value): + update(item, f"{path}/{index}") + else: + _update_field(digest, repr(value).encode()) + + chained = getattr(optimizer, "chained_optimizers", None) + if chained is not None: + optimizers = [getattr(item, "optimizer", item) for item in chained] + else: + optimizers = [getattr(optimizer, "optimizer", optimizer)] + for optimizer_index, inner_optimizer in enumerate(optimizers): + param_groups = getattr(inner_optimizer, "param_groups", None) + state = getattr(inner_optimizer, "state", None) + if param_groups is None or state is None: + update(inner_optimizer.state_dict(), f"optimizer/{optimizer_index}/state_dict") + continue + for group_index, group in enumerate(param_groups): + update( + {key: value for key, value in group.items() if key != "params"}, + f"optimizer/{optimizer_index}/group/{group_index}/options", + ) + for parameter_index, parameter in enumerate(group["params"]): + path = f"optimizer/{optimizer_index}/group/{group_index}/parameter/{parameter_index}" + update(parameter, f"{path}/master") + update(state.get(parameter, {}), f"{path}/state") + return digest.hexdigest() + + +def reference_probe_sha256( + pair_ids: Sequence[int], + branch_is_chosen: Sequence[bool], + tokens: Sequence[Sequence[int] | torch.Tensor], + loss_masks: Sequence[Sequence[int] | torch.Tensor], + ref_log_probs: Sequence[Sequence[float] | torch.Tensor], +) -> str: + """Hash the exact completion-only frozen-reference probe output.""" + size = len(pair_ids) + if any(len(values) != size for values in (branch_is_chosen, tokens, loss_masks, ref_log_probs)): + raise ValueError("reference probe fields must be branch aligned") + digest = hashlib.sha256() + for index in range(size): + _update_field(digest, str(int(pair_ids[index])).encode()) + _update_field(digest, b"chosen" if bool(branch_is_chosen[index]) else b"rejected") + token_tensor = torch.as_tensor(tokens[index], dtype=torch.int64) + mask_tensor = torch.as_tensor(loss_masks[index], dtype=torch.bool) + logp_tensor = torch.as_tensor(ref_log_probs[index], dtype=torch.float32) + if logp_tensor.shape != mask_tensor.shape: + raise ValueError("reference probe log-probability/mask shape mismatch") + _update_field(digest, _tensor_bytes(token_tensor)) + _update_field(digest, _tensor_bytes(mask_tensor)) + masked_log_probs = logp_tensor * mask_tensor.to(device=logp_tensor.device, dtype=torch.float32) + _update_field(digest, _tensor_bytes(masked_log_probs)) + return digest.hexdigest() + + +def reference_identity_path(checkpoint_root: str | os.PathLike[str], iteration: int) -> Path: + root = Path(checkpoint_root) + iteration_dir = root if root.name == f"iter_{iteration:07d}" else root / f"iter_{iteration:07d}" + return iteration_dir / REFERENCE_IDENTITY_FILENAME + + +def write_reference_identity(path: Path, identity: DPOReferenceIdentity) -> None: + """Atomically write a reference identity sidecar.""" + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(f"{path.suffix}.tmp.{os.getpid()}") + temporary.write_text(json.dumps(identity.to_dict(), indent=2, sort_keys=True) + "\n", encoding="utf-8") + os.replace(temporary, path) + + +def read_reference_identity(path: Path) -> DPOReferenceIdentity: + if not path.is_file(): + raise FileNotFoundError(f"DPO reference identity sidecar is missing: {path}") + identity = DPOReferenceIdentity.from_dict(json.loads(path.read_text(encoding="utf-8"))) + if identity.schema_version != 1: + raise ValueError(f"unsupported DPO reference identity schema: {identity.schema_version}") + return identity + + +__all__ = [ + "DPOReferenceIdentity", + "REFERENCE_IDENTITY_FILENAME", + "REFERENCE_LOADER_MODE", + "canonical_optimizer_sha256", + "canonical_tensor_sha256", + "read_reference_identity", + "reference_identity_path", + "reference_probe_sha256", + "resolve_dpo_reference_checkpoint", + "write_reference_identity", +] diff --git a/relax/components/actor.py b/relax/components/actor.py index c915cf90f..a8d98b587 100644 --- a/relax/components/actor.py +++ b/relax/components/actor.py @@ -14,7 +14,7 @@ from relax.components.base import Base from relax.distributed.coordination import PeerStepBarrier, RolloutOffloadBarrier from relax.distributed.ray.placement_group import allocate_train_group -from relax.engine.sft.runtime import is_sft_mode, sft_partition_id, sft_task_name +from relax.engine.sft.runtime import is_preference_mode, is_sft_mode, sft_partition_id, sft_task_name from relax.utils.async_utils import run from relax.utils.opd.opd_utils import set_managed_opd_teacher_on_train_group @@ -78,7 +78,11 @@ def __init__( self.actor_model.async_init( config, role=self.role, - with_ref=config.kl_coef != 0 or config.use_kl_loss, + with_ref=( + config.kl_coef != 0 + or config.use_kl_loss + or (is_preference_mode(config) and config.sft_objective == "dpo" and not config.dpo_reference_free) + ), with_opd_teacher=self.config.opd_teacher_load, ) ) diff --git a/relax/components/sft.py b/relax/components/sft.py index 15ca2b15c..53860ed7b 100644 --- a/relax/components/sft.py +++ b/relax/components/sft.py @@ -34,8 +34,14 @@ from transformers import AutoConfig, AutoTokenizer from relax.components.base import Base +from relax.engine.sft.dataset.preference import ( + PreferenceStreamingDataset, + ProcessedPreferencePair, + pack_preference_pairs_for_tq, +) from relax.engine.sft.dataset.streaming import ProcessedSample, SFTStreamingDataset, pack_samples_for_tq from relax.engine.sft.debug_print import print_first_sample +from relax.engine.sft.runtime import is_preference_mode from relax.utils.data.processor_pool import ProcessorPool from relax.utils.misc import load_function from relax.utils.s3_model_loader import prepare_model_maybe_update_args @@ -98,11 +104,12 @@ def _init_data_pipeline(self) -> None: return prepare_model_maybe_update_args(self.config, completeness="metadata") self._tokenizer = AutoTokenizer.from_pretrained(self.config.hf_checkpoint, trust_remote_code=True) - try: - self._processor_pool = ProcessorPool(self.config.hf_checkpoint, pool_size=None, trust_remote_code=True) - except Exception as exc: - self._logger.warning(f"Could not init ProcessorPool ({exc}); multimodal samples will fail at push.") - self._processor_pool = None + if not is_preference_mode(self.config): + try: + self._processor_pool = ProcessorPool(self.config.hf_checkpoint, pool_size=None, trust_remote_code=True) + except Exception as exc: + self._logger.warning(f"Could not init ProcessorPool ({exc}); multimodal samples will fail at push.") + self._processor_pool = None pad_token_ids = _resolve_pad_token_ids_from_config(self.config.hf_checkpoint) self._logger.info(f"Resolved multimodal pad token ids from model config: {sorted(pad_token_ids)}") @@ -127,7 +134,27 @@ def _init_data_pipeline(self) -> None: self._logger.info(f"SFT invalid multimodal strategy: {invalid_multimodal_strategy}") dataset_cls = _load_custom_dataset_class(getattr(self.config, "custom_dataset_class_path", None)) - if dataset_cls is None: + if is_preference_mode(self.config): + self._dataset = PreferenceStreamingDataset( + path=self.config.prompt_data, + tokenizer=self._tokenizer, + prompt_key=self.config.input_key, + chosen_key=self.config.preference_chosen_key, + rejected_key=self.config.preference_rejected_key, + pair_id_key=self.config.preference_pair_id_key, + metadata_key=self.config.metadata_key, + max_length=self.config.preference_max_length, + max_completion_length=self.config.preference_max_completion_length, + pair_capacity=capacity, + seed=seed, + prefetch_max_cached=prefetch_buffer_size, + prefetch_chunk_size=prefetch_chunk_size, + prefetch_num_workers=prefetch_num_workers, + apply_chat_template_kwargs=getattr(self.config, "apply_chat_template_kwargs", None), + expected_chat_template_sha256=self.config.preference_chat_template_sha256, + require_no_generation_marker=self.config.preference_require_no_generation_marker, + ) + elif dataset_cls is None: self._dataset = SFTStreamingDataset( path=self.config.prompt_data, tokenizer=self._tokenizer, @@ -300,10 +327,16 @@ async def _wait_for_buffer_capacity(self) -> None: wait_count += 1 await asyncio.sleep(1) - def _maybe_print_first_sample(self, samples: list[ProcessedSample]) -> None: + def _maybe_print_first_sample(self, samples: list[ProcessedSample] | list[ProcessedPreferencePair]) -> None: if self.step != 0 or not samples: return s = samples[0] + if isinstance(s, ProcessedPreferencePair): + self._logger.info( + f"First preference pair: pair_id={s.pair_id!r}, " + f"chosen_length={s.chosen_total_length}, rejected_length={s.rejected_total_length}" + ) + return try: print_first_sample( step=self.step, @@ -333,12 +366,23 @@ async def _produce_one_step(self) -> None: "consumer requires a full global batch. Check invalid-multimodal and oversize skip warnings." ) self._maybe_print_first_sample(samples) - backend_batch = pack_samples_for_tq(samples, force_multimodal_field=self.config.multimodal_keys is not None) - assert backend_batch is not None - await self.data_system_client.async_put( - data=dict_to_tensordict(backend_batch, batch_size=len(backend_batch["tokens"])), - partition_id=f"sft_{self.step}", - ) + if is_preference_mode(self.config): + backend_batch, custom_meta = pack_preference_pairs_for_tq(samples) + batch_size = len(backend_batch["pair_ids"]) + await self.data_system_client.async_put( + data=dict_to_tensordict(backend_batch, batch_size=batch_size), + partition_id=f"sft_{self.step}", + custom_meta=custom_meta, + ) + else: + backend_batch = pack_samples_for_tq( + samples, force_multimodal_field=self.config.multimodal_keys is not None + ) + assert backend_batch is not None + await self.data_system_client.async_put( + data=dict_to_tensordict(backend_batch, batch_size=len(backend_batch["tokens"])), + partition_id=f"sft_{self.step}", + ) if crossed_epoch: self._logger.info( f"SFT step {self.step}: epoch boundary crossed (epoch={self._dataset.index_manager.current_epoch})" diff --git a/relax/engine/sft/bootstrap.py b/relax/engine/sft/bootstrap.py index 5871bde28..b22c7e55e 100644 --- a/relax/engine/sft/bootstrap.py +++ b/relax/engine/sft/bootstrap.py @@ -58,11 +58,21 @@ def resolve_sft_num_rollout(config: Namespace) -> None: # Lazy import: pulling streaming dataset at module load would drag heavy # multimodal deps into every controller import. - from relax.engine.sft.dataset.streaming import SFTStreamingDataset + if getattr(config, "sft_objective", "causal_lm") == "dpo": + from relax.engine.sft.dataset.preference import PreferenceStreamingDataset + + sizing_dataset = PreferenceStreamingDataset( + path=config.prompt_data, + pair_id_key=config.preference_pair_id_key, + prefetch_max_cached=0, + ) + else: + from relax.engine.sft.dataset.streaming import SFTStreamingDataset + + sizing_dataset = SFTStreamingDataset(path=config.prompt_data, prefetch_max_cached=0) # Sized-only construction: no tokenizer/processor needed because we never # call get_batch — we just need len() to derive num_rollout. - sizing_dataset = SFTStreamingDataset(path=config.prompt_data, prefetch_max_cached=0) dataset_size = len(sizing_dataset) num_per_epoch = dataset_size // config.rollout_batch_size assert num_per_epoch > 0, f"SFT dataset size {dataset_size} < rollout_batch_size {config.rollout_batch_size}" diff --git a/relax/engine/sft/dataset/preference.py b/relax/engine/sft/dataset/preference.py new file mode 100644 index 000000000..3f154a4f4 --- /dev/null +++ b/relax/engine/sft/dataset/preference.py @@ -0,0 +1,534 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Streaming chosen/rejected dataset for offline preference objectives.""" + +import hashlib +import threading +from collections import Counter +from dataclasses import dataclass +from typing import Any + +import torch + +from relax.engine.sft.dataset.chat_template import ( + HAS_GENERATION_MARKER, + _resolve_sft_template_kwargs, + render_with_loss_mask, +) +from relax.engine.sft.dataset.sample import CanonicalMessage, CanonicalSample +from relax.engine.sft.dataset.streaming import _build_reader +from relax.utils.data.streaming_dataset import IndexManager, PrefetchBuffer +from relax.utils.logging_utils import get_logger + + +logger = get_logger(__name__) + + +class PreferenceDataError(ValueError): + """Stable, classified preference-row rejection.""" + + def __init__(self, reason_code: str, message: str, *, source_idx: int, pair_id: str | None = None) -> None: + super().__init__(message) + self.reason_code = reason_code + self.source_idx = source_idx + self.pair_id = pair_id + + +class _PreferenceRowError(ValueError): + """Internal row rejection that carries its stable reason code.""" + + def __init__(self, reason_code: str, message: str) -> None: + super().__init__(message) + self.reason_code = reason_code + + +def _classify_preference_error(error: BaseException) -> str: + # Prefer the reason code attached at the raise site; message matching is + # only a fallback for errors raised outside this module (e.g. tokenizer). + reason_code = getattr(error, "reason_code", None) + if isinstance(reason_code, str) and reason_code: + return reason_code + message = str(error).lower() + if "post-truncation identical" in message: + return "post_truncation" + if "must not be identical" in message: + return "identical" + if "prompt tokens differ" in message or "strict common prefix" in message: + return "prompt_mismatch" + if "empty completion" in message or "no supervised tokens" in message: + return "empty_completion" + if "capacity" in message or "exceeds max_length" in message: + return "oversize" + return "schema" + + +@dataclass(frozen=True) +class PreferencePair: + """Canonical text-only preference pair before tokenization.""" + + pair_id: str + prompt: list[CanonicalMessage] + chosen: CanonicalMessage + rejected: CanonicalMessage + metadata: dict[str, Any] + + +@dataclass(frozen=True) +class ProcessedPreferencePair: + """Tokenized pair kept atomic until after DP assignment.""" + + pair_id: str + chosen_tokens: torch.Tensor + rejected_tokens: torch.Tensor + chosen_loss_mask: torch.Tensor + rejected_loss_mask: torch.Tensor + chosen_total_length: int + rejected_total_length: int + chosen_prompt_length: int + rejected_prompt_length: int + chosen_completion_length: int + rejected_completion_length: int + source_idx: int + + @property + def pair_total_length(self) -> int: + return self.chosen_total_length + self.rejected_total_length + + +def _message_from_raw(raw: Any, *, learn: bool, field: str) -> CanonicalMessage: + if not isinstance(raw, dict): + raise _PreferenceRowError("schema", f"preference {field} must be a message object") + role = raw.get("role") + content = raw.get("content") + if role is None or content is None: + raise _PreferenceRowError("schema", f"preference {field} message requires role and content") + if not isinstance(content, str): + raise _PreferenceRowError("schema", f"preference {field} must be pure text") + return CanonicalMessage(role=role, content=content, learn=learn, tool_calls=raw.get("tool_calls")) + + +def _normalize_pair_row( + row: dict[str, Any], + *, + row_index: int, + prompt_key: str, + chosen_key: str, + rejected_key: str, + pair_id_key: str, + metadata_key: str, + source_name: str, +) -> PreferencePair: + pair_id = row.get(pair_id_key) + if not isinstance(pair_id, str) or not pair_id: + raise _PreferenceRowError("schema", f"preference row requires a non-empty {pair_id_key}") + chosen_raw = row.get(chosen_key) + rejected_raw = row.get(rejected_key) + prompt_raw = row.get(prompt_key) + + if prompt_raw is None: + if not isinstance(chosen_raw, list) or not isinstance(rejected_raw, list): + raise _PreferenceRowError("schema", "implicit preference rows require chosen/rejected message lists") + prefix_length = 0 + for chosen_message, rejected_message in zip(chosen_raw, rejected_raw, strict=False): + if chosen_message != rejected_message: + break + prefix_length += 1 + chosen_suffix = chosen_raw[prefix_length:] + rejected_suffix = rejected_raw[prefix_length:] + if len(chosen_suffix) != 1 or len(rejected_suffix) != 1: + raise _PreferenceRowError( + "prompt_mismatch", + "implicit preference rows require one assistant message after the strict common prefix", + ) + prompt_raw = chosen_raw[:prefix_length] + chosen_raw = chosen_suffix[0] + rejected_raw = rejected_suffix[0] + elif not isinstance(prompt_raw, list): + raise _PreferenceRowError("schema", f"preference {prompt_key} must be a message list") + + prompt = [_message_from_raw(message, learn=False, field=prompt_key) for message in prompt_raw] + chosen = _message_from_raw(chosen_raw, learn=True, field=chosen_key) + rejected = _message_from_raw(rejected_raw, learn=True, field=rejected_key) + if chosen.role != "assistant" or rejected.role != "assistant": + raise _PreferenceRowError("schema", "preference chosen/rejected messages must have role assistant") + if chosen.content == rejected.content: + raise _PreferenceRowError("identical", "preference chosen/rejected responses must not be identical") + + metadata = row.get(metadata_key) or {} + if not isinstance(metadata, dict): + raise _PreferenceRowError("schema", f"preference {metadata_key} must be an object") + metadata = dict(metadata) + metadata.update({"source_dataset": source_name, "row_index": row_index, "pair_id": pair_id}) + return PreferencePair(pair_id=pair_id, prompt=prompt, chosen=chosen, rejected=rejected, metadata=metadata) + + +def _split_branch( + tokens: torch.Tensor, mask: torch.Tensor, *, pair_id: str, branch: str +) -> tuple[torch.Tensor, torch.Tensor]: + if tokens.ndim != 1 or mask.ndim != 1 or tokens.shape != mask.shape: + raise _PreferenceRowError( + "schema", f"preference pair {pair_id!r} {branch} tokens/mask must be aligned one-dimensional tensors" + ) + supervised = torch.nonzero(mask, as_tuple=False).flatten() + if supervised.numel() == 0: + raise _PreferenceRowError( + "empty_completion", f"preference pair {pair_id!r} {branch} completion has no supervised tokens" + ) + first = int(supervised[0]) + if not bool(mask[first:].to(dtype=torch.bool).all()): + raise _PreferenceRowError( + "schema", f"preference pair {pair_id!r} {branch} completion mask must be one contiguous suffix" + ) + return tokens[:first], tokens[first:] + + +def _truncate_pair( + prompt: torch.Tensor, + chosen_completion: torch.Tensor, + rejected_completion: torch.Tensor, + *, + pair_id: str, + max_length: int, + max_completion_length: int, + pair_capacity: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + chosen_completion = chosen_completion[:max_completion_length] + rejected_completion = rejected_completion[:max_completion_length] + if chosen_completion.numel() == 0 or rejected_completion.numel() == 0: + raise _PreferenceRowError( + "empty_completion", f"preference pair {pair_id!r} has an empty completion after truncation" + ) + prompt_budget = max_length - max(chosen_completion.numel(), rejected_completion.numel()) + if prompt_budget < 0: + raise _PreferenceRowError( + "oversize", f"preference pair {pair_id!r} completion exceeds max_length={max_length}" + ) + prompt = prompt[-prompt_budget:] if prompt_budget else prompt[:0] + + def total() -> int: + return 2 * prompt.numel() + chosen_completion.numel() + rejected_completion.numel() + + while total() > pair_capacity: + if chosen_completion.numel() >= rejected_completion.numel() and chosen_completion.numel() > 1: + chosen_completion = chosen_completion[:-1] + elif rejected_completion.numel() > 1: + rejected_completion = rejected_completion[:-1] + elif prompt.numel() > 0: + prompt = prompt[1:] + else: + raise _PreferenceRowError( + "oversize", + f"preference pair {pair_id!r} cannot fit pair capacity {pair_capacity} while retaining both completions", + ) + return prompt.contiguous(), chosen_completion.contiguous(), rejected_completion.contiguous() + + +class PreferenceStreamingDataset: + """Lazy text-only preference dataset with deterministic epoch shuffling.""" + + def __init__( + self, + path: str | list[str] | tuple[str, ...], + *, + tokenizer=None, + prompt_key: str = "prompt", + chosen_key: str = "chosen", + rejected_key: str = "rejected", + pair_id_key: str = "prompt_id", + metadata_key: str = "metadata", + source_name: str = "preference_data", + max_length: int = 1024, + max_completion_length: int = 512, + pair_capacity: int | None = None, + seed: int = 42, + prefetch_max_cached: int = 256, + prefetch_chunk_size: int = 32, + prefetch_num_workers: int = 4, + apply_chat_template_kwargs: dict | None = None, + expected_chat_template_sha256: str | None = None, + require_no_generation_marker: bool = False, + ) -> None: + if max_length <= 0 or max_completion_length <= 0: + raise ValueError("preference length limits must be positive") + self.reader = _build_reader(path) + self.index_manager = IndexManager(len(self.reader), seed=seed) + self.tokenizer = tokenizer + self.prompt_key = prompt_key + self.chosen_key = chosen_key + self.rejected_key = rejected_key + self.pair_id_key = pair_id_key + self.metadata_key = metadata_key + self.source_name = source_name + self.max_length = max_length + self.max_completion_length = max_completion_length + self.pair_capacity = pair_capacity or (2 * max_length) + self.apply_chat_template_kwargs = apply_chat_template_kwargs + self.expected_chat_template_sha256 = expected_chat_template_sha256 + self.require_no_generation_marker = require_no_generation_marker + self._rejection_counts: Counter[str] = Counter() + self._rejection_records: list[dict[str, Any]] = [] + self._rejection_lock = threading.Lock() + self._template_contract_validated = False + self._template_contract_lock = threading.Lock() + self._validate_unique_pair_ids() + self._first_error: BaseException | None = None + self._error_lock = threading.Lock() + self._prefetch: PrefetchBuffer | None = None + if prefetch_max_cached > 0: + self._prefetch = PrefetchBuffer( + process_fn=self._process_one_safe, + chunk_size=prefetch_chunk_size, + max_cached=prefetch_max_cached, + num_workers=prefetch_num_workers, + ) + + def __len__(self) -> int: + return len(self.reader) + + def _validate_unique_pair_ids(self) -> None: + seen: set[str] = set() + for index in range(len(self.reader)): + pair_id = self.reader[index].get(self.pair_id_key) + if not isinstance(pair_id, str) or not pair_id: + message = f"preference row {index} requires a non-empty {self.pair_id_key}" + with self._rejection_lock: + self._rejection_counts["schema"] += 1 + self._rejection_records.append( + {"source_idx": index, "pair_id": None, "reason_code": "schema", "message": message} + ) + raise PreferenceDataError("schema", message, source_idx=index) + if pair_id in seen: + message = f"duplicate preference pair ID {pair_id!r} at row {index}" + with self._rejection_lock: + self._rejection_counts["schema"] += 1 + self._rejection_records.append( + {"source_idx": index, "pair_id": pair_id, "reason_code": "schema", "message": message} + ) + raise PreferenceDataError("schema", message, source_idx=index, pair_id=pair_id) + seen.add(pair_id) + + def shuffle(self, epoch_id: int, position: int = 0) -> None: + self.index_manager.shuffle(epoch_id) + if position: + self.index_manager.position = min(position, self.index_manager.total_size) + if self._prefetch is not None and self.index_manager.indices is not None: + remaining = self.index_manager.indices[self.index_manager.position :] + self._prefetch.set_index_order(list(remaining)) + + def stop(self) -> None: + if self._prefetch is not None: + self._prefetch.stop() + + def get_canonical_pair(self, idx: int) -> PreferencePair: + return _normalize_pair_row( + self.reader[idx], + row_index=idx, + prompt_key=self.prompt_key, + chosen_key=self.chosen_key, + rejected_key=self.rejected_key, + pair_id_key=self.pair_id_key, + metadata_key=self.metadata_key, + source_name=self.source_name, + ) + + def get_processed_pair(self, idx: int) -> ProcessedPreferencePair: + try: + return self._get_processed_pair(idx) + except PreferenceDataError: + raise + except Exception as exc: + pair_id = None + try: + raw_pair_id = self.reader[idx].get(self.pair_id_key) + pair_id = raw_pair_id if isinstance(raw_pair_id, str) else None + except Exception: + pass + reason_code = _classify_preference_error(exc) + error = PreferenceDataError(reason_code, str(exc), source_idx=idx, pair_id=pair_id) + with self._rejection_lock: + self._rejection_counts[reason_code] += 1 + self._rejection_records.append( + {"source_idx": idx, "pair_id": pair_id, "reason_code": reason_code, "message": str(exc)} + ) + logger.error( + "Rejected preference pair source_idx=%s pair_id=%r reason_code=%s counts=%s", + idx, + pair_id, + reason_code, + dict(self.rejection_counts), + ) + raise error from exc + + def _get_processed_pair(self, idx: int) -> ProcessedPreferencePair: + if self.tokenizer is None: + raise RuntimeError("PreferenceStreamingDataset requires a tokenizer for processing") + pair = self.get_canonical_pair(idx) + chosen_sample = CanonicalSample(messages=[*pair.prompt, pair.chosen], metadata=dict(pair.metadata), tools=None) + rejected_sample = CanonicalSample( + messages=[*pair.prompt, pair.rejected], metadata=dict(pair.metadata), tools=None + ) + self._validate_template_contract(chosen_sample) + chosen_tokens, chosen_mask = render_with_loss_mask( + chosen_sample, + tokenizer=self.tokenizer, + apply_chat_template_kwargs=self.apply_chat_template_kwargs, + ) + rejected_tokens, rejected_mask = render_with_loss_mask( + rejected_sample, + tokenizer=self.tokenizer, + apply_chat_template_kwargs=self.apply_chat_template_kwargs, + ) + chosen_prompt, chosen_completion = _split_branch( + chosen_tokens, chosen_mask, pair_id=pair.pair_id, branch="chosen" + ) + rejected_prompt, rejected_completion = _split_branch( + rejected_tokens, rejected_mask, pair_id=pair.pair_id, branch="rejected" + ) + if not torch.equal(chosen_prompt, rejected_prompt): + raise _PreferenceRowError( + "prompt_mismatch", f"preference pair {pair.pair_id!r} chosen/rejected prompt tokens differ" + ) + prompt, chosen_completion, rejected_completion = _truncate_pair( + chosen_prompt, + chosen_completion, + rejected_completion, + pair_id=pair.pair_id, + max_length=self.max_length, + max_completion_length=self.max_completion_length, + pair_capacity=self.pair_capacity, + ) + chosen_tokens = torch.cat((prompt, chosen_completion)) + rejected_tokens = torch.cat((prompt, rejected_completion)) + if torch.equal(chosen_tokens, rejected_tokens) or torch.equal(chosen_completion, rejected_completion): + raise _PreferenceRowError( + "post_truncation", f"preference pair {pair.pair_id!r} is post-truncation identical" + ) + chosen_mask = torch.cat((torch.zeros_like(prompt), torch.ones_like(chosen_completion))) + rejected_mask = torch.cat((torch.zeros_like(prompt), torch.ones_like(rejected_completion))) + return ProcessedPreferencePair( + pair_id=pair.pair_id, + chosen_tokens=chosen_tokens, + rejected_tokens=rejected_tokens, + chosen_loss_mask=chosen_mask, + rejected_loss_mask=rejected_mask, + chosen_total_length=chosen_tokens.numel(), + rejected_total_length=rejected_tokens.numel(), + chosen_prompt_length=prompt.numel(), + rejected_prompt_length=prompt.numel(), + chosen_completion_length=chosen_completion.numel(), + rejected_completion_length=rejected_completion.numel(), + source_idx=idx, + ) + + @property + def rejection_counts(self) -> dict[str, int]: + """Return a thread-safe snapshot of classified rejection counts.""" + with self._rejection_lock: + return dict(self._rejection_counts) + + @property + def rejection_records(self) -> list[dict[str, Any]]: + """Return row IDs and stable reason codes for evidence manifests.""" + with self._rejection_lock: + return [dict(record) for record in self._rejection_records] + + def _validate_template_contract(self, sample: CanonicalSample) -> None: + if self._template_contract_validated: + return + with self._template_contract_lock: + if self._template_contract_validated: + return + resolved = _resolve_sft_template_kwargs( + sample, + tokenizer=self.tokenizer, + apply_chat_template_kwargs=self.apply_chat_template_kwargs, + ) + template = resolved.template or "" + digest = hashlib.sha256(template.encode()).hexdigest() + if self.expected_chat_template_sha256 is not None and digest != self.expected_chat_template_sha256: + raise ValueError( + "preference chat template SHA-256 mismatch: " + f"expected={self.expected_chat_template_sha256}, actual={digest}" + ) + if self.require_no_generation_marker and HAS_GENERATION_MARKER(template): + raise ValueError("preference recipe requires a chat template without {% generation %} markers") + self._template_contract_validated = True + + def _process_one_safe(self, idx: int) -> ProcessedPreferencePair | None: + try: + return self.get_processed_pair(idx) + except Exception as exc: + with self._error_lock: + if self._first_error is None: + self._first_error = exc + logger.exception(f"PreferenceStreamingDataset: failed to process pair idx={idx}") + return None + + def _raise_if_failed(self) -> None: + with self._error_lock: + error = self._first_error + if error is not None: + raise error + + def get_batch(self, n: int) -> tuple[list[ProcessedPreferencePair], bool]: + if n <= 0: + raise ValueError(f"batch size must be positive, got {n}") + self._raise_if_failed() + pairs: list[ProcessedPreferencePair] = [] + crossed_epoch = False + attempts = 0 + max_attempts = max(n * 10, 32) + while len(pairs) < n and attempts < max_attempts: + indices, crossed = self.index_manager.get_next_indices(1) + crossed_epoch = crossed_epoch or crossed + attempts += 1 + index = indices[0] + pair = self._prefetch.get(index) if self._prefetch is not None else self.get_processed_pair(index) + if pair is None: + self._raise_if_failed() + else: + pairs.append(pair) + if len(pairs) != n: + raise RuntimeError(f"preference dataset returned a partial batch: expected {n}, got {len(pairs)}") + return pairs, crossed_epoch + + async def get_batch_async(self, n: int) -> tuple[list[ProcessedPreferencePair], bool]: + return self.get_batch(n) + + def get_batch_in_order(self, start: int, n: int) -> list[ProcessedPreferencePair]: + return [self.get_processed_pair(index) for index in range(start, min(start + n, len(self.reader)))] + + +def pack_preference_pairs_for_tq( + pairs: list[ProcessedPreferencePair], +) -> tuple[dict[str, list[Any]], list[dict[str, int]]]: + """Pack atomic pair rows and matching TransferQueue length metadata.""" + if not pairs: + raise ValueError("preference pair batch must not be empty") + encoded_pair_ids = [ + int.from_bytes(hashlib.sha256(pair.pair_id.encode()).digest()[:8], "big") >> 1 for pair in pairs + ] + if len(set(encoded_pair_ids)) != len(encoded_pair_ids): + raise ValueError("preference pair ID hash collision within batch") + batch: dict[str, list[Any]] = { + "pair_ids": encoded_pair_ids, + "chosen_tokens": [pair.chosen_tokens.tolist() for pair in pairs], + "rejected_tokens": [pair.rejected_tokens.tolist() for pair in pairs], + "chosen_loss_masks": [pair.chosen_loss_mask.tolist() for pair in pairs], + "rejected_loss_masks": [pair.rejected_loss_mask.tolist() for pair in pairs], + "chosen_total_lengths": [pair.chosen_total_length for pair in pairs], + "rejected_total_lengths": [pair.rejected_total_length for pair in pairs], + } + custom_meta = [{"total_lengths": pair.pair_total_length} for pair in pairs] + if len(custom_meta) != len(pairs): + raise RuntimeError("preference custom metadata is not row-aligned") + return batch, custom_meta + + +__all__ = [ + "PreferenceDataError", + "PreferencePair", + "PreferenceStreamingDataset", + "ProcessedPreferencePair", + "pack_preference_pairs_for_tq", +] diff --git a/relax/engine/sft/runtime.py b/relax/engine/sft/runtime.py index d5c1a51cc..792a8b95c 100644 --- a/relax/engine/sft/runtime.py +++ b/relax/engine/sft/runtime.py @@ -7,6 +7,7 @@ here keeps the dispatchers in those files to one-line calls. """ +import math from argparse import Namespace @@ -19,6 +20,88 @@ def is_sft_mode(args: Namespace) -> bool: return getattr(args, "loss_type", None) == "sft" +def sft_objective(args: Namespace) -> str: + """Return the offline objective while preserving causal-LM defaults.""" + return getattr(args, "sft_objective", "causal_lm") + + +def is_preference_mode(args: Namespace) -> bool: + return is_sft_mode(args) and sft_objective(args) == "dpo" + + +def validate_preference_args(args: Namespace) -> None: + """Reject unsupported preference configurations before Serve starts.""" + if not is_preference_mode(args): + return + if getattr(args, "custom_dataset_class_path", None): + raise ValueError("preference objectives do not support --custom-dataset-class") + if getattr(args, "multimodal_keys", None) is not None: + raise ValueError("preference objectives v1 support pure text only") + if int(getattr(args, "n_samples_per_prompt", 1)) != 1: + raise ValueError("preference objectives require --n-samples-per-prompt 1") + topology = { + "tensor_model_parallel_size": int(getattr(args, "tensor_model_parallel_size", 1) or 1), + "pipeline_model_parallel_size": int(getattr(args, "pipeline_model_parallel_size", 1) or 1), + "context_parallel_size": int(getattr(args, "context_parallel_size", 1) or 1), + } + invalid = {name: size for name, size in topology.items() if size != 1} + if invalid: + raise ValueError(f"preference objectives v1 require TP=CP=PP=1, got {invalid}") + if getattr(args, "dynamic_context_parallel", False): + raise ValueError("preference objectives v1 do not support dynamic context parallelism") + if getattr(args, "qkv_format", "thd") != "thd": + raise ValueError("preference objectives v1 require --qkv-format thd") + if getattr(args, "fully_async", False) or getattr(args, "hybrid", False): + raise ValueError("preference objectives v1 support synchronous SFT topology only") + if not getattr(args, "use_gloo_process_groups", False): + raise ValueError("preference objectives require --use-gloo-process-groups for DP iterator control data") + if getattr(args, "sft_chunked_logits", False) or getattr(args, "enable_mtp_training", False): + raise ValueError("preference objectives v1 do not support SFT chunked logits or MTP") + if getattr(args, "calculate_per_token_loss", False): + raise ValueError("preference objectives use pair reduction and reject --calculate-per-token-loss") + if int(getattr(args, "lora_rank", 0) or 0) > 0: + raise ValueError("preference objectives v1 do not support LoRA") + if ( + float(getattr(args, "hidden_dropout", 0.0) or 0.0) != 0.0 + or float(getattr(args, "attention_dropout", 0.0) or 0.0) != 0.0 + ): + raise ValueError("preference objectives require hidden and attention dropout to be 0.0") + if getattr(args, "sft_predict_interval", None) is not None: + raise ValueError("preference objectives do not use SFT generation prediction") + max_length = int(getattr(args, "preference_max_length", 0) or 0) + max_completion_length = int(getattr(args, "preference_max_completion_length", 0) or 0) + if max_length <= 0 or max_completion_length <= 0: + raise ValueError("preference length limits must be positive") + if max_completion_length > max_length: + raise ValueError("--preference-max-completion-length must not exceed --preference-max-length") + seq_length = int(getattr(args, "seq_length", max_length) or max_length) + if max_length > seq_length: + raise ValueError("--preference-max-length must not exceed --seq-length") + if bool(getattr(args, "eval_prompt_data", None)) or getattr(args, "eval_size", None) is not None: + raise ValueError("DPO held-out evaluation is delivered by the follow-up reward-modeling PR") + beta = float(getattr(args, "dpo_beta", 0.1)) + if not math.isfinite(beta) or beta <= 0: + raise ValueError(f"--dpo-beta must be finite and positive, got {beta}") + likelihood_temperature = float(getattr(args, "rollout_temperature", 1.0)) + if not math.isfinite(likelihood_temperature) or likelihood_temperature != 1.0: + raise ValueError( + "preference objectives require --rollout-temperature 1.0 so sampling temperature does not scale " + "policy/reference likelihood logits" + ) + if getattr(args, "ref_load", None) is not None: + raise ValueError( + "preference objectives do not use --ref-load: standard DPO snapshots the frozen reference " + "from the pinned --dpo-reference-repository/--dpo-reference-revision snapshot" + ) + if not getattr(args, "dpo_reference_free", False) and getattr(args, "ref_update_interval", None) is not None: + raise ValueError("standard DPO requires a frozen reference and rejects --ref-update-interval") + if not getattr(args, "dpo_reference_free", False) and not getattr(args, "enable_weights_backuper", False): + raise ValueError("standard DPO requires --enable-weights-backuper for actor/ref snapshots") + if not getattr(args, "dpo_reference_free", False): + if not getattr(args, "dpo_reference_repository", None) or not getattr(args, "dpo_reference_revision", None): + raise ValueError("standard DPO requires --dpo-reference-repository and --dpo-reference-revision") + + def sft_partition_id(args: Namespace, step: int) -> str: return f"sft_{step}" if is_sft_mode(args) else f"train_{step}" diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index f7d1e710e..ece532f89 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -578,6 +578,28 @@ def add_train_arguments(parser): ) # ---- SFT / Predict ---- + parser.add_argument( + "--sft-objective", + choices=["causal_lm", "dpo"], + default="causal_lm", + help="Offline objective under --loss-type sft. Defaults to the existing causal-LM behavior.", + ) + parser.add_argument("--preference-chosen-key", type=str, default="chosen") + parser.add_argument("--preference-rejected-key", type=str, default="rejected") + parser.add_argument("--preference-pair-id-key", type=str, default="prompt_id") + parser.add_argument("--preference-max-length", type=int, default=1024) + parser.add_argument("--preference-max-completion-length", type=int, default=512) + parser.add_argument("--preference-chat-template-sha256", type=str, default=None) + parser.add_argument("--preference-require-no-generation-marker", action="store_true", default=False) + parser.add_argument("--dpo-beta", type=float, default=0.1) + parser.add_argument("--dpo-reference-repository", type=str, default=None) + parser.add_argument("--dpo-reference-revision", type=str, default=None) + parser.add_argument( + "--dpo-reference-free", + action=argparse.BooleanOptionalAction, + default=False, + help="Use explicit reference-free logistic DPO instead of a frozen reference checkpoint.", + ) parser.add_argument( "--custom-dataset-class", "--custom-dataset-class-path", @@ -3315,6 +3337,9 @@ def slime_validate_args(args): if not args.balance_data: logger.info("--loss-type sft: auto-enabling --balance-data for DP-balanced batching.") args.balance_data = True + from relax.engine.sft.runtime import validate_preference_args + + validate_preference_args(args) args.use_critic = args.advantage_estimator == "ppo" # Synchronous PPO has no producer for diff --git a/relax/utils/data/stream_dataloader.py b/relax/utils/data/stream_dataloader.py index f8b35115d..c51f2bd87 100644 --- a/relax/utils/data/stream_dataloader.py +++ b/relax/utils/data/stream_dataloader.py @@ -946,6 +946,15 @@ def _fetch_once() -> list: def post_process_rollout_data(args, rollout_data): # move tokens/loss_masks to GPU in-place as a list of tensors (downstream # code in this module expects lists of sequence tensors for packing) + if "tokens" not in rollout_data and "chosen_tokens" in rollout_data: + # Preference rows stay pair-atomic in TransferQueue. Expand them only + # after sampling, before the generic sequence post-processing below. + from relax.backends.megatron.data import expand_preference_rollout_data + + expanded = expand_preference_rollout_data(rollout_data) + rollout_data.clear() + rollout_data.update(expanded) + from relax.backends.megatron.cp_utils import maybe_padded_total_lengths, slice_log_prob_with_cp cuda_dev = device_utils.make_current_torch_device() diff --git a/relax/utils/training/data_fields.py b/relax/utils/training/data_fields.py index ebd6479af..101a5d7ba 100644 --- a/relax/utils/training/data_fields.py +++ b/relax/utils/training/data_fields.py @@ -27,6 +27,16 @@ def build_data_fields(args: Namespace, *, consumer: str = "actor") -> list[str]: ``actor``); other algorithms ignore it and receive the base rollout fields. """ if getattr(args, "loss_type", None) == "sft": + if getattr(args, "sft_objective", "causal_lm") == "dpo": + return [ + "pair_ids", + "chosen_tokens", + "rejected_tokens", + "chosen_loss_masks", + "rejected_loss_masks", + "chosen_total_lengths", + "rejected_total_lengths", + ] fields = ["tokens", "total_lengths", "response_lengths", "loss_masks"] if args.multimodal_keys is not None: fields.append("multimodal_train_inputs") diff --git a/relax/utils/training/preference_utils.py b/relax/utils/training/preference_utils.py new file mode 100644 index 000000000..b0df6d939 --- /dev/null +++ b/relax/utils/training/preference_utils.py @@ -0,0 +1,156 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Pure helpers for pair-aware DPO training.""" + +from collections.abc import Sequence + +import torch +import torch.nn.functional as F + + +def require_tensor_condition(condition: torch.Tensor, message: str) -> None: + """Raise on CPU immediately and enqueue a device-side assertion on CUDA.""" + if condition.device.type == "cuda": + torch._assert_async(condition, message) + elif not bool(condition): + raise ValueError(message) + + +def _validate_same_shape(name: str, *values: torch.Tensor) -> None: + if not values: + raise ValueError(f"{name} requires at least one tensor") + expected = values[0].shape + if any(value.shape != expected for value in values[1:]): + shapes = [tuple(value.shape) for value in values] + raise ValueError(f"{name} tensors must have identical shapes, got {shapes}") + + +def build_causal_lm_labels(tokens: torch.Tensor, raw_loss_mask: torch.Tensor) -> torch.Tensor: + """Build next-token labels from an unshifted completion-token mask.""" + if tokens.ndim != 1 or raw_loss_mask.ndim != 1: + raise ValueError("tokens and raw_loss_mask must be one-dimensional") + if tokens.shape != raw_loss_mask.shape: + raise ValueError(f"tokens/raw_loss_mask shape mismatch: {tuple(tokens.shape)} vs {tuple(raw_loss_mask.shape)}") + labels = torch.full_like(tokens, -100) + if tokens.numel() > 1: + supervised = raw_loss_mask[1:].to(dtype=torch.bool) + labels[:-1][supervised] = tokens[1:][supervised] + return labels + + +def dpo_pair_loss( + policy_chosen: torch.Tensor, + policy_rejected: torch.Tensor, + *, + reference_chosen: torch.Tensor | None = None, + reference_rejected: torch.Tensor | None = None, + beta: float = 0.1, + reference_free: bool = False, +) -> torch.Tensor: + """Return unreduced sigmoid-DPO loss, one value per preference pair.""" + if beta <= 0: + raise ValueError(f"DPO beta must be positive, got {beta}") + _validate_same_shape("policy log-probabilities", policy_chosen, policy_rejected) + policy_logratio = policy_chosen - policy_rejected + if reference_free: + if reference_chosen is not None or reference_rejected is not None: + raise ValueError("reference-free DPO must not receive reference log-probabilities") + reference_logratio = torch.zeros_like(policy_logratio) + else: + if reference_chosen is None or reference_rejected is None: + raise ValueError("standard DPO requires chosen and rejected reference log-probabilities") + _validate_same_shape( + "reference log-probabilities", + policy_chosen, + reference_chosen, + reference_rejected, + ) + reference_logratio = reference_chosen - reference_rejected + logits = beta * (policy_logratio - reference_logratio) + require_tensor_condition(torch.isfinite(logits).all(), "DPO logits must contain only finite values") + return -F.logsigmoid(logits) + + +def build_preference_pair_indices( + branch_pair_ids: Sequence[int], branch_is_chosen: Sequence[bool] +) -> tuple[list[int], list[int]]: + """Return chosen/rejected branch indices grouped by stable pair + identity.""" + if len(branch_pair_ids) != len(branch_is_chosen): + raise ValueError( + "preference pair identity fields must be branch aligned: " + f"{len(branch_pair_ids)} vs {len(branch_is_chosen)}" + ) + if not branch_pair_ids: + raise ValueError("DPO micro-batch must contain at least one preference pair") + + pairs: dict[int, dict[bool, int]] = {} + order: list[int] = [] + for index, (raw_pair_id, raw_is_chosen) in enumerate(zip(branch_pair_ids, branch_is_chosen, strict=True)): + pair_id = int(raw_pair_id) + is_chosen = bool(raw_is_chosen) + if pair_id not in pairs: + pairs[pair_id] = {} + order.append(pair_id) + if is_chosen in pairs[pair_id]: + branch = "chosen" if is_chosen else "rejected" + raise ValueError(f"preference pair {pair_id!r} contains duplicate {branch} branches") + pairs[pair_id][is_chosen] = index + + chosen_indices: list[int] = [] + rejected_indices: list[int] = [] + for pair_id in order: + pair = pairs[pair_id] + if set(pair) != {False, True}: + raise ValueError(f"preference pair {pair_id!r} must contain exactly one chosen and one rejected branch") + chosen_indices.append(pair[True]) + rejected_indices.append(pair[False]) + return chosen_indices, rejected_indices + + +def pack_preference_pair_indices( + costs: Sequence[int], + pair_ids: Sequence[str], + *, + capacity: int, +) -> list[list[int]]: + """Deterministic capacity-aware first-fit-decreasing pair packing.""" + if capacity <= 0: + raise ValueError(f"capacity must be positive, got {capacity}") + if len(costs) != len(pair_ids): + raise ValueError(f"costs/pair_ids length mismatch: {len(costs)} vs {len(pair_ids)}") + normalized_costs = [int(cost) for cost in costs] + for pair_id, cost in zip(pair_ids, normalized_costs, strict=True): + if cost <= 0: + raise ValueError(f"pair {pair_id!r} has non-positive cost {cost}") + if cost > capacity: + raise ValueError(f"oversize preference pair {pair_id!r} has cost {cost}, capacity={capacity}") + + order = sorted(range(len(normalized_costs)), key=lambda index: (-normalized_costs[index], str(pair_ids[index]))) + bins: list[list[int]] = [] + bin_costs: list[int] = [] + for index in order: + cost = normalized_costs[index] + for bin_index, bin_cost in enumerate(bin_costs): + if bin_cost + cost <= capacity: + bins[bin_index].append(index) + bin_costs[bin_index] += cost + break + else: + bins.append([index]) + bin_costs.append(cost) + + if sorted(index for group in bins for index in group) != list(range(len(normalized_costs))): + raise RuntimeError("preference pair packer lost or duplicated pair indices") + if any(sum(normalized_costs[index] for index in group) > capacity for group in bins): + raise RuntimeError("preference pair packer produced an over-capacity micro-batch") + return bins + + +__all__ = [ + "build_preference_pair_indices", + "build_causal_lm_labels", + "dpo_pair_loss", + "pack_preference_pair_indices", + "require_tensor_condition", +] diff --git a/scripts/data/prepare_ultrafeedback_preferences.py b/scripts/data/prepare_ultrafeedback_preferences.py new file mode 100644 index 000000000..12c13f304 --- /dev/null +++ b/scripts/data/prepare_ultrafeedback_preferences.py @@ -0,0 +1,219 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Create deterministic Task 31 train/eval preference subsets and a provenance +manifest.""" + +import argparse +import hashlib +import json +from pathlib import Path +from typing import Any + +from relax.engine.sft.dataset.sample import VALID_ROLES + + +DATASET_ID = "HuggingFaceH4/ultrafeedback_binarized" +DATASET_REVISION = "3949bf5f8c17c394422ccfab0c31ea9c20bdeb85" +ORDER_NAMESPACE = "task31-ultrafeedback-v1:" +REJECTION_REASON_CODES = ( + "schema", + "identical", + "post_truncation", + "prompt_mismatch", + "empty_completion", + "oversize", +) + + +def _reason_code(error: BaseException) -> str: + message = str(error).lower() + if "identical" in message: + return "identical" + if "strict shared prompt" in message: + return "prompt_mismatch" + return "schema" + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _validate_row(row: dict[str, Any], *, split: str, index: int) -> str: + prompt_id = row.get("prompt_id") + if not isinstance(prompt_id, str) or not prompt_id: + raise ValueError(f"{split}[{index}] has no non-empty prompt_id") + chosen = row.get("chosen") + rejected = row.get("rejected") + if not isinstance(chosen, list) or not isinstance(rejected, list) or not chosen or not rejected: + raise ValueError(f"{split}[{index}] {prompt_id} has invalid chosen/rejected messages") + for branch_name, messages in (("chosen", chosen), ("rejected", rejected)): + for message_index, message in enumerate(messages): + if not isinstance(message, dict): + raise ValueError( + f"{split}[{index}] {prompt_id} {branch_name}[{message_index}] must be a message object" + ) + role = message.get("role") + content = message.get("content") + if role not in VALID_ROLES: + raise ValueError( + f"{split}[{index}] {prompt_id} {branch_name}[{message_index}] has invalid role {role!r}" + ) + if not isinstance(content, str): + raise ValueError( + f"{split}[{index}] {prompt_id} {branch_name}[{message_index}] content must be a string" + ) + if chosen == rejected: + raise ValueError(f"{split}[{index}] {prompt_id} has identical chosen/rejected branches") + if chosen[-1].get("role") != "assistant" or rejected[-1].get("role") != "assistant": + raise ValueError(f"{split}[{index}] {prompt_id} must end both branches with assistant") + if chosen[:-1] != rejected[:-1]: + raise ValueError(f"{split}[{index}] {prompt_id} does not have a strict shared prompt") + return prompt_id + + +def _select(dataset, *, split: str, count: int) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + seen: set[str] = set() + candidates: list[tuple[str, str, dict[str, Any]]] = [] + rejected: list[dict[str, Any]] = [] + for index, row in enumerate(dataset): + if not isinstance(row, dict): + rejected.append( + { + "prompt_id": None, + "source_index": index, + "reason_code": "schema", + "reason": f"{split}[{index}] must be an object", + } + ) + continue + prompt_id = row.get("prompt_id") + if not isinstance(prompt_id, str) or not prompt_id: + raise ValueError(f"{split}[{index}] has no non-empty prompt_id") + if prompt_id in seen: + rejected.append( + { + "prompt_id": prompt_id, + "source_index": index, + "reason_code": "schema", + "reason": f"duplicate prompt_id in {split}; retained first source occurrence", + } + ) + continue + seen.add(prompt_id) + order_key = hashlib.sha256(f"{ORDER_NAMESPACE}{prompt_id}".encode()).hexdigest() + candidates.append((order_key, prompt_id, {**row, "_source_index": index})) + if len(candidates) < count: + raise ValueError(f"{split} contains {len(candidates)} valid rows, need {count}") + candidates.sort(key=lambda item: (item[0], item[1])) + selected: list[dict[str, Any]] = [] + for _, prompt_id, row in candidates: + source_index = row.pop("_source_index") + try: + _validate_row(row, split=split, index=source_index) + except ValueError as exc: + rejected.append( + { + "prompt_id": prompt_id, + "source_index": source_index, + "reason_code": _reason_code(exc), + "reason": str(exc), + } + ) + continue + selected.append( + { + "prompt_id": prompt_id, + "chosen": row["chosen"], + "rejected": row["rejected"], + "metadata": { + "source_split": split, + "score_chosen": row.get("score_chosen"), + "score_rejected": row.get("score_rejected"), + }, + } + ) + if len(selected) == count: + break + if len(selected) != count: + raise ValueError(f"{split} contains only {len(selected)} valid rows after deterministic validation") + return selected, rejected + + +def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None: + with path.open("w", encoding="utf-8", newline="\n") as stream: + for row in rows: + stream.write(json.dumps(row, ensure_ascii=False, sort_keys=True, separators=(",", ":")) + "\n") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--train-size", type=int, default=4096) + parser.add_argument("--eval-size", type=int, default=512) + args = parser.parse_args() + if args.train_size <= 0 or args.eval_size <= 0: + parser.error("subset sizes must be positive") + + from datasets import Dataset, load_dataset + + train_source = load_dataset(DATASET_ID, revision=DATASET_REVISION, split="train_prefs") + eval_source = load_dataset(DATASET_ID, revision=DATASET_REVISION, split="test_prefs") + train_rows, train_rejections = _select(train_source, split="train_prefs", count=args.train_size) + eval_rows, eval_rejections = _select(eval_source, split="test_prefs", count=args.eval_size) + train_ids = {row["prompt_id"] for row in train_rows} + eval_ids = {row["prompt_id"] for row in eval_rows} + overlap = train_ids & eval_ids + if overlap: + raise ValueError(f"train/eval prompt_id overlap: {sorted(overlap)[:5]}") + + args.output_dir.mkdir(parents=True, exist_ok=True) + outputs: dict[str, dict[str, Any]] = {} + for name, rows in (("train", train_rows), ("eval", eval_rows)): + jsonl_path = args.output_dir / f"ultrafeedback_{name}.jsonl" + parquet_path = args.output_dir / f"ultrafeedback_{name}.parquet" + _write_jsonl(jsonl_path, rows) + Dataset.from_list(rows).to_parquet(parquet_path) + outputs[name] = { + "count": len(rows), + "prompt_ids": [row["prompt_id"] for row in rows], + "jsonl": {"path": jsonl_path.name, "sha256": _sha256(jsonl_path)}, + "parquet": {"path": parquet_path.name, "sha256": _sha256(parquet_path)}, + } + + manifest = { + "schema_version": 1, + "source": {"dataset": DATASET_ID, "revision": DATASET_REVISION}, + "selection": { + "algorithm": f'sha256("{ORDER_NAMESPACE}" + prompt_id), then first valid rows', + "overlap_count": 0, + "rejections": { + "train": train_rejections, + "eval": eval_rejections, + }, + "rejection_counts": { + split: { + reason_code: sum(item["reason_code"] == reason_code for item in rejections) + for reason_code in REJECTION_REASON_CODES + } + for split, rejections in (("train", train_rejections), ("eval", eval_rejections)) + }, + }, + "schema": { + "prompt_id": "string", + "chosen": "list<{role:string,content:string}>", + "rejected": "list<{role:string,content:string}>", + "metadata": "object", + }, + "outputs": outputs, + } + manifest_path = args.output_dir / "manifest.json" + manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") + print(json.dumps({"manifest": str(manifest_path), "sha256": _sha256(manifest_path)}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/scripts/training/dpo/run-qwen3-0.6B-ultrafeedback-1xgpu.sh b/scripts/training/dpo/run-qwen3-0.6B-ultrafeedback-1xgpu.sh new file mode 100644 index 000000000..2ccd1d2c7 --- /dev/null +++ b/scripts/training/dpo/run-qwen3-0.6B-ultrafeedback-1xgpu.sh @@ -0,0 +1,65 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +set -eo pipefail +set -x + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +if [ -z "${RELAX_ENTRYPOINT_MODE:-}" ]; then + source "${SCRIPT_DIR}/../../entrypoint/local.sh" +fi +source "${MODEL_CONFIG_DIR}/qwen3-0.6B.sh" + +MODEL_REVISION=c1899de289a04d12100db370d81485cdf75e47ca +HF_CHECKPOINT="${HF_CHECKPOINT:-${MODEL_DIR}/Qwen3-0.6B-${MODEL_REVISION}}" +PROMPT_DATA="${PROMPT_DATA:?set PROMPT_DATA to the Task 31 train JSONL or Parquet}" +SAVE_DIR="${SAVE_DIR:-${SCRIPT_DIR}/../../../checkpoints/task31-dpo}" +EXP_NAME="${EXP_NAME:-qwen3-0.6b-ultrafeedback-dpo-gpu1}" +now=$(date "+%Y-%m-%d-%H:%M:%S") + +mkdir -p log "${SAVE_DIR}" +ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://127.0.0.1:8265" \ + ${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 -m relax.entrypoints.train \ + --resource '{"sft":[1,0],"actor":[1,1]}' \ + --loss-type sft \ + --sft-objective dpo \ + --dpo-beta 0.1 \ + --dpo-reference-repository Qwen/Qwen3-0.6B \ + --dpo-reference-revision "${MODEL_REVISION}" \ + --prompt-data "${PROMPT_DATA}" \ + --input-key prompt \ + --preference-pair-id-key prompt_id \ + --preference-max-length 1024 \ + --preference-max-completion-length 512 \ + --preference-chat-template-sha256 56965952fc78cd889bcd1864d70e85271861eef93385410b879c0c4c2d40564d \ + --preference-require-no-generation-marker \ + --hf-checkpoint "${HF_CHECKPOINT}" \ + --megatron-to-hf-mode bridge \ + --enable-weights-backuper \ + --save "${SAVE_DIR}/${EXP_NAME}" \ + --load "${SAVE_DIR}/${EXP_NAME}" \ + --save-interval "${SAVE_INTERVAL:-50}" \ + --num-rollout "${NUM_ROLLOUT:-200}" \ + --global-batch-size "${GLOBAL_BATCH_SIZE:-32}" \ + --use-dynamic-batch-size \ + --use-gloo-process-groups \ + --max-tokens-per-gpu "${MAX_TOKENS_PER_GPU:-8192}" \ + --tensor-model-parallel-size 1 \ + --pipeline-model-parallel-size 1 \ + --context-parallel-size 1 \ + --optimizer adam \ + --lr "${LR:-5e-7}" \ + --lr-decay-style cosine \ + --min-lr 0 \ + ${OVERRIDE_OPT_PARAM_SCHEDULER:+--override-opt-param-scheduler} \ + --weight-decay 0.0 \ + --clip-grad 1.0 \ + --attention-dropout 0.0 \ + --hidden-dropout 0.0 \ + --attention-backend flash \ + --no-rope-fusion \ + --colocate \ + "${MODEL_ARGS[@]}" 2>&1 | tee "log/${EXP_NAME}-${now}.log" diff --git a/tests/backends/megatron/test_dpo_loss.py b/tests/backends/megatron/test_dpo_loss.py new file mode 100644 index 000000000..7f7045b1f --- /dev/null +++ b/tests/backends/megatron/test_dpo_loss.py @@ -0,0 +1,139 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Production DPO loss regression tests.""" + +from argparse import Namespace + +import pytest +import torch +import torch.nn.functional as F + + +try: + from relax.backends.megatron import loss as loss_module +except Exception as exc: + pytest.skip(f"relax.backends.megatron unavailable: {exc}", allow_module_level=True) + + +def _args(*, reference_free: bool = False, beta: float = 0.2) -> Namespace: + return Namespace(dpo_reference_free=reference_free, dpo_beta=beta) + + +def _run(monkeypatch, policy_values, *, order=None, reference_free=False, ref_values=None, num_samples=2): + if order is None: + order = [0, 1, 2, 3] + pair_ids = [10, 10, 20, 20] + is_chosen = [True, False, True, False] + policy = [torch.as_tensor(policy_values[index]).reshape(1) for index in order] + reference = None if ref_values is None else [torch.as_tensor(ref_values[index]).reshape(1) for index in order] + monkeypatch.setattr( + loss_module, "get_log_probs_and_entropy", lambda *args, **kwargs: (None, {"log_probs": policy}) + ) + logits = torch.ones(1, requires_grad=True) + batch = { + "response_lengths": [1] * 4, + "unconcat_tokens": [torch.ones(1, dtype=torch.long)] * 4, + "total_lengths": [1] * 4, + "loss_masks": [torch.ones(1)] * 4, + "preference_branch_pair_ids": [pair_ids[index] for index in order], + "preference_is_chosen": [is_chosen[index] for index in order], + "ref_log_probs": reference, + "num_samples": num_samples, + } + return loss_module.dpo_loss_function(_args(reference_free=reference_free), batch, logits, lambda value: value) + + +def test_production_dpo_loss_matches_independent_reference_and_gradients(monkeypatch): + policy = torch.tensor([-1.0, -2.0, -0.5, -0.75], requires_grad=True) + reference = torch.tensor([-1.2, -1.7, -0.4, -0.8]) + actual, metrics = _run(monkeypatch, list(policy.unbind()), ref_values=list(reference.unbind())) + expected = -F.logsigmoid(0.2 * ((policy[0::2] - policy[1::2]) - (reference[0::2] - reference[1::2]))).sum() + torch.testing.assert_close(actual, expected) + actual.backward() + actual_grad = policy.grad.clone() + policy.grad = None + expected.backward() + torch.testing.assert_close(actual_grad, policy.grad) + assert { + "dpo/logps_chosen", + "dpo/logps_rejected", + "dpo/ref_logps_chosen", + "dpo/ref_logps_rejected", + "dpo/tie_rate", + "dpo/tie_aware_accuracy", + }.issubset(metrics) + + +def test_pair_identity_restores_reordered_micro_batch(monkeypatch): + policy = [-1.0, -2.0, -0.5, -0.75] + reference = [-1.2, -1.7, -0.4, -0.8] + baseline, baseline_metrics = _run(monkeypatch, policy, ref_values=reference) + reordered, reordered_metrics = _run(monkeypatch, policy, ref_values=reference, order=[3, 0, 2, 1]) + torch.testing.assert_close(reordered, baseline) + for key in baseline_metrics: + torch.testing.assert_close(reordered_metrics[key], baseline_metrics[key]) + + +def test_tie_metrics_are_epsilon_aware(monkeypatch): + _, metrics = _run( + monkeypatch, + [-1.0, -2.0, -0.5, -0.75], + ref_values=[-1.0, -2.0, -0.5, -0.75], + ) + assert metrics["dpo/strict_accuracy"].item() == 0 + assert metrics["dpo/tie_rate"].item() == 2 + assert metrics["dpo/tie_aware_accuracy"].item() == 1 + + +def test_reference_free_partition_and_num_samples_do_not_change_pair_sum(monkeypatch): + policy = [-1.0, -2.0, -0.5, -0.75] + first, _ = _run(monkeypatch, policy, reference_free=True, num_samples=1) + second, _ = _run(monkeypatch, policy, reference_free=True, num_samples=999) + torch.testing.assert_close(first, second) + pair_losses = -F.logsigmoid(0.2 * (torch.tensor(policy)[0::2] - torch.tensor(policy)[1::2])) + torch.testing.assert_close(first, pair_losses[:1].sum() + pair_losses[1:].sum()) + + +@pytest.mark.parametrize( + ("pair_ids", "chosen", "match"), + [ + ([1, 1, 2, 2], [True, True, True, False], "duplicate chosen"), + ([1, 2, 2, 3], [True, True, False, False], "exactly one"), + ], +) +def test_production_dpo_rejects_invalid_pair_identity(monkeypatch, pair_ids, chosen, match): + monkeypatch.setattr( + loss_module, + "get_log_probs_and_entropy", + lambda *args, **kwargs: (None, {"log_probs": [torch.zeros(1) for _ in pair_ids]}), + ) + batch = { + "response_lengths": [1] * len(pair_ids), + "unconcat_tokens": [torch.ones(1, dtype=torch.long)] * len(pair_ids), + "total_lengths": [1] * len(pair_ids), + "loss_masks": [torch.ones(1)] * len(pair_ids), + "preference_branch_pair_ids": pair_ids, + "preference_is_chosen": chosen, + "ref_log_probs": [torch.zeros(1) for _ in pair_ids], + } + with pytest.raises(ValueError, match=match): + loss_module.dpo_loss_function(_args(), batch, torch.ones(1), lambda value: value) + + +def test_production_dpo_rejects_empty_completion_mask(monkeypatch): + monkeypatch.setattr( + loss_module, + "get_log_probs_and_entropy", + lambda *args, **kwargs: (None, {"log_probs": [torch.zeros(1), torch.zeros(1)]}), + ) + batch = { + "response_lengths": [1, 1], + "unconcat_tokens": [torch.ones(1, dtype=torch.long)] * 2, + "total_lengths": [1, 1], + "loss_masks": [torch.zeros(1), torch.ones(1)], + "preference_branch_pair_ids": [1, 1], + "preference_is_chosen": [True, False], + "ref_log_probs": [torch.zeros(1), torch.zeros(1)], + } + with pytest.raises(ValueError, match="at least one supervised token"): + loss_module.dpo_loss_function(_args(), batch, torch.ones(1), lambda value: value) diff --git a/tests/backends/megatron/test_dpo_reference_integrity.py b/tests/backends/megatron/test_dpo_reference_integrity.py new file mode 100644 index 000000000..60d71097f --- /dev/null +++ b/tests/backends/megatron/test_dpo_reference_integrity.py @@ -0,0 +1,374 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Frozen-reference checksum, probe, optimizer and sidecar tests.""" + +import hashlib +import json +import os +import sys +import types +from argparse import Namespace + +import pytest +import torch + +from relax.backends.megatron.reference_integrity import ( + DPOReferenceIdentity, + canonical_optimizer_sha256, + canonical_tensor_sha256, + read_reference_identity, + reference_probe_sha256, + resolve_dpo_reference_checkpoint, + write_reference_identity, +) + + +def test_megatron_resume_detection_ignores_fresh_output_directory(tmp_path): + pytest.importorskip("megatron.training.checkpointing") + from relax.backends.megatron.checkpoint import is_megatron_checkpoint + + output = tmp_path / "run" + output.mkdir() + (output / "transformer_config.json").write_text("{}", encoding="utf-8") + assert not is_megatron_checkpoint(output) + (output / "latest_checkpointed_iteration.txt").write_text("1", encoding="utf-8") + assert is_megatron_checkpoint(output) + assert is_megatron_checkpoint(tmp_path / "iter_0000001") + + +def test_canonical_tensor_digest_is_order_stable_and_byte_sensitive(): + first = canonical_tensor_sha256([("b", torch.tensor([2.0])), ("a", torch.tensor([1.0]))]) + reordered = canonical_tensor_sha256([("a", torch.tensor([1.0])), ("b", torch.tensor([2.0]))]) + changed = canonical_tensor_sha256([("a", torch.tensor([1.0])), ("b", torch.tensor([3.0]))]) + assert first == reordered + assert first != changed + assert first != canonical_tensor_sha256([("a", torch.tensor([1], dtype=torch.int64)), ("b", torch.tensor([2.0]))]) + + +def test_optimizer_digest_detects_master_or_state_changes(): + parameter = torch.nn.Parameter(torch.tensor([1.0])) + optimizer = torch.optim.Adam([parameter], lr=0.1) + (parameter.square().sum()).backward() + optimizer.step() + baseline = canonical_optimizer_sha256(optimizer) + master_value = parameter.detach().clone() + parameter.data.add_(1) + master_changed = canonical_optimizer_sha256(optimizer) + assert master_changed != baseline + parameter.data.copy_(master_value) + optimizer.state[parameter]["exp_avg"].add_(1) + assert canonical_optimizer_sha256(optimizer) != baseline + + +def test_probe_digest_covers_identity_tokens_masks_and_fp32_logprobs(): + args = ([1, 1], [True, False], [[1, 2], [1, 3]], [[0, 1], [0, 1]]) + baseline = reference_probe_sha256(*args, [[0.0, -1.0], [0.0, -2.0]]) + assert baseline != reference_probe_sha256(*args, [[0.0, -1.0], [0.0, -2.1]]) + assert baseline != reference_probe_sha256([2, 2], *args[1:], [[0.0, -1.0], [0.0, -2.0]]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for cross-device probe coverage") +def test_probe_digest_accepts_gpu_logprobs_with_cpu_manifest(): + cpu_digest = reference_probe_sha256([1], [True], [[1, 2]], [[0, 1]], [[0.0, -1.0]]) + gpu_digest = reference_probe_sha256([1], [True], [[1, 2]], [[0, 1]], [torch.tensor([0.0, -1.0], device="cuda")]) + assert gpu_digest == cpu_digest + + +def test_reference_identity_sidecar_is_required_and_rejects_schema_damage(tmp_path): + path = tmp_path / "relax_dpo_reference.json" + with pytest.raises(FileNotFoundError): + read_reference_identity(path) + identity = DPOReferenceIdentity(1, "repo", "revision", "loader", "a" * 64, "b" * 64) + write_reference_identity(path, identity) + assert read_reference_identity(path) == identity + payload = identity.to_dict() + payload["schema_version"] = 99 + path.write_text(json.dumps(payload), encoding="utf-8") + with pytest.raises(ValueError, match="unsupported"): + read_reference_identity(path) + + +def _write_local_download_metadata(checkpoint, filename, revision): + source = checkpoint / filename + metadata = checkpoint / ".cache" / "huggingface" / "download" / f"{filename}.metadata" + metadata.parent.mkdir(parents=True, exist_ok=True) + content = source.read_bytes() + if source.suffix == ".safetensors": + etag = hashlib.sha256(content).hexdigest() + else: + etag = hashlib.sha1(f"blob {len(content)}\0".encode() + content).hexdigest() + metadata.write_text(f"{revision}\n{etag}\n{source.stat().st_mtime}\n", encoding="utf-8") + + +def test_resolve_dpo_reference_checkpoint_uses_the_pinned_configured_local_snapshot(monkeypatch, tmp_path): + revision = "a" * 40 + snapshot = tmp_path / f"Qwen3-0.6B-{revision}" + snapshot.mkdir() + (snapshot / "config.json").write_text("{}", encoding="utf-8") + (snapshot / "model.safetensors").write_bytes(b"weights") + _write_local_download_metadata(snapshot, "config.json", revision) + _write_local_download_metadata(snapshot, "model.safetensors", revision) + observed = {} + + def snapshot_download(**kwargs): + observed.update(kwargs) + return str(snapshot) + + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + types.SimpleNamespace(snapshot_download=snapshot_download), + ) + assert resolve_dpo_reference_checkpoint("org/model", revision, str(snapshot)) == str(snapshot.resolve()) + assert observed == { + "repo_id": "org/model", + "revision": revision, + "local_dir": str(snapshot.resolve()), + "local_files_only": True, + } + + +def test_resolve_dpo_reference_checkpoint_rejects_missing_local_snapshot(monkeypatch, tmp_path): + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + + def snapshot_download(**_kwargs): + raise OSError("not cached") + + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + types.SimpleNamespace(snapshot_download=snapshot_download), + ) + with pytest.raises(RuntimeError, match="hf download org/model --revision"): + resolve_dpo_reference_checkpoint("org/model", "a" * 40, str(checkpoint)) + + +def test_resolve_dpo_reference_checkpoint_rejects_different_resolved_directory(monkeypatch, tmp_path): + checkpoint = tmp_path / "checkpoint" + other = tmp_path / "other" + checkpoint.mkdir() + other.mkdir() + (checkpoint / "config.json").write_text("{}", encoding="utf-8") + (other / "config.json").write_text("{}", encoding="utf-8") + + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + types.SimpleNamespace(snapshot_download=lambda **_kwargs: str(other)), + ) + with pytest.raises(RuntimeError, match="different from --hf-checkpoint"): + resolve_dpo_reference_checkpoint("org/model", "a" * 40, str(checkpoint)) + + +def test_resolve_dpo_reference_checkpoint_rejects_missing_pinned_local_metadata(monkeypatch, tmp_path): + revision = "a" * 40 + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text("{}", encoding="utf-8") + (checkpoint / "model.safetensors").write_bytes(b"weights") + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + types.SimpleNamespace(snapshot_download=lambda **_kwargs: str(checkpoint)), + ) + with pytest.raises(RuntimeError, match="missing valid Hugging Face local-dir metadata"): + resolve_dpo_reference_checkpoint("org/model", revision, str(checkpoint)) + + +def test_resolve_dpo_reference_checkpoint_rejects_mismatched_file_metadata(monkeypatch, tmp_path): + revision = "a" * 40 + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text("{}", encoding="utf-8") + _write_local_download_metadata(checkpoint, "config.json", "c" * 40) + (checkpoint / "model.safetensors").write_bytes(b"weights") + _write_local_download_metadata(checkpoint, "model.safetensors", revision) + + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + types.SimpleNamespace(snapshot_download=lambda **_kwargs: str(checkpoint)), + ) + with pytest.raises(RuntimeError, match="metadata does not match the pinned revision"): + resolve_dpo_reference_checkpoint("org/model", revision, str(checkpoint)) + + +def test_resolve_dpo_reference_checkpoint_rejects_replaced_file_with_restored_mtime(monkeypatch, tmp_path): + revision = "a" * 40 + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + config = checkpoint / "config.json" + config.write_text("{}", encoding="utf-8") + _write_local_download_metadata(checkpoint, "config.json", revision) + weights = checkpoint / "model.safetensors" + weights.write_bytes(b"weights") + original_stat = weights.stat() + _write_local_download_metadata(checkpoint, "model.safetensors", revision) + weights.write_bytes(b"changed") + os.utime(weights, ns=(original_stat.st_atime_ns, original_stat.st_mtime_ns)) + + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + types.SimpleNamespace(snapshot_download=lambda **_kwargs: str(checkpoint)), + ) + with pytest.raises(RuntimeError, match="contents do not match its Hugging Face ETag"): + resolve_dpo_reference_checkpoint("org/model", revision, str(checkpoint)) + + +def test_resolve_dpo_reference_checkpoint_rejects_snapshot_without_supported_weights(monkeypatch, tmp_path): + revision = "a" * 40 + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text("{}", encoding="utf-8") + _write_local_download_metadata(checkpoint, "config.json", revision) + + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + types.SimpleNamespace(snapshot_download=lambda **_kwargs: str(checkpoint)), + ) + with pytest.raises(RuntimeError, match="no supported model weights or index"): + resolve_dpo_reference_checkpoint("org/model", revision, str(checkpoint)) + + +def test_resolve_dpo_reference_checkpoint_accepts_complete_safetensors_index(monkeypatch, tmp_path): + revision = "a" * 40 + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text("{}", encoding="utf-8") + index_name = "model.safetensors.index.json" + shard_names = ("model-00001-of-00002.safetensors", "model-00002-of-00002.safetensors") + weight_map = {f"layer.{index}.weight": shard for index, shard in enumerate(shard_names)} + (checkpoint / index_name).write_text(json.dumps({"weight_map": weight_map}), encoding="utf-8") + for shard_name in shard_names: + (checkpoint / shard_name).write_bytes(shard_name.encode()) + for filename in ("config.json", index_name, *shard_names): + _write_local_download_metadata(checkpoint, filename, revision) + + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + types.SimpleNamespace(snapshot_download=lambda **_kwargs: str(checkpoint)), + ) + assert resolve_dpo_reference_checkpoint("org/model", revision, str(checkpoint)) == str(checkpoint.resolve()) + + +def test_resolve_dpo_reference_checkpoint_rejects_missing_indexed_shard_and_metadata(monkeypatch, tmp_path): + revision = "a" * 40 + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text("{}", encoding="utf-8") + index_name = "model.safetensors.index.json" + shard_name = "model-00001-of-00001.safetensors" + (checkpoint / index_name).write_text(json.dumps({"weight_map": {"model.weight": shard_name}}), encoding="utf-8") + (checkpoint / shard_name).write_bytes(b"weights") + for filename in ("config.json", index_name, shard_name): + _write_local_download_metadata(checkpoint, filename, revision) + (checkpoint / shard_name).unlink() + (checkpoint / ".cache" / "huggingface" / "download" / f"{shard_name}.metadata").unlink() + + monkeypatch.setitem( + sys.modules, + "huggingface_hub", + types.SimpleNamespace(snapshot_download=lambda **_kwargs: str(checkpoint)), + ) + with pytest.raises(RuntimeError, match="weight index points to a missing shard"): + resolve_dpo_reference_checkpoint("org/model", revision, str(checkpoint)) + + +def test_resume_probe_materializes_manifest_lists_as_tensors(monkeypatch): + try: + from relax.backends.megatron import actor as actor_module + except Exception as exc: + pytest.skip(f"Megatron actor unavailable: {exc}") + + instance = object.__new__(actor_module.MegatronTrainRayActor) + instance.args = Namespace() + instance.model = [object()] + instance._dpo_reference_probe_verified = False + instance._expected_dpo_reference_identity = DPOReferenceIdentity( + 1, + "repo", + "revision", + "loader", + "a" * 64, + "b" * 64, + { + "pair_ids": [1, 1], + "branch_is_chosen": [True, False], + "tokens": [[1, 2], [1, 3]], + "loss_masks": [[False, True], [False, True]], + "total_lengths": [2, 2], + "response_lengths": [1, 1], + }, + ) + + def inspect_probe_data(_args, _model, probe_data): + assert all(torch.is_tensor(value) and value.dtype == torch.long for value in probe_data["tokens"]) + assert all(torch.is_tensor(value) and value.dtype == torch.bool for value in probe_data["loss_masks"]) + raise RuntimeError("probe inspected") + + monkeypatch.setattr(actor_module, "get_data_iterator", inspect_probe_data) + monkeypatch.setattr(actor_module.mpu, "get_data_parallel_world_size", lambda **_kwargs: 1) + monkeypatch.setattr(actor_module.device_utils, "make_current_torch_device", lambda: torch.device("cpu")) + with pytest.raises(RuntimeError, match="probe inspected"): + instance._replay_dpo_reference_probe() + + +def test_loader_half_write_failure_restores_actor_and_keeps_optimizer(monkeypatch): + try: + from relax.backends.megatron import actor as actor_module + except Exception as exc: + pytest.skip(f"Megatron actor unavailable: {exc}") + + parameter = torch.nn.Parameter(torch.tensor([1.0])) + optimizer = torch.optim.Adam([parameter], lr=0.1) + + class _Backuper: + def __init__(self): + self.values = {"actor": {"weight": parameter.detach().clone()}} + + @property + def backup_tags(self): + return list(self.values) + + def restore(self, tag): + parameter.data.copy_(self.values[tag]["weight"]) + + def backup(self, tag): + self.values[tag] = {"weight": parameter.detach().clone()} + + def get(self, tag): + return self.values[tag] + + instance = object.__new__(actor_module.MegatronTrainRayActor) + instance.args = Namespace( + load="checkpoint", + no_load_optim=False, + no_load_rng=False, + finetune=False, + megatron_to_hf_mode="bridge", + dpo_reference_repository="repo", + dpo_reference_revision="revision", + ) + instance.model = [object()] + instance.optimizer = optimizer + instance.weights_backuper = _Backuper() + instance._active_model_tag = "actor" + instance._expected_dpo_reference_identity = None + monkeypatch.setattr(actor_module.device_utils, "maybe_backend_process_on_model_switch", lambda: None) + + def fail_after_half_write(*args, **kwargs): + parameter.data.fill_(99) + raise RuntimeError("injected loader failure") + + monkeypatch.setattr(actor_module, "load_checkpoint", fail_after_half_write) + before = canonical_optimizer_sha256(optimizer) + with pytest.raises(RuntimeError, match="injected loader failure"): + instance._rebuild_dpo_reference("hf-path") + assert parameter.item() == 1.0 + assert instance._active_model_tag == "actor" + assert "ref" not in instance.weights_backuper.backup_tags + assert canonical_optimizer_sha256(optimizer) == before diff --git a/tests/backends/megatron/test_preference_batching.py b/tests/backends/megatron/test_preference_batching.py new file mode 100644 index 000000000..15c027b0b --- /dev/null +++ b/tests/backends/megatron/test_preference_batching.py @@ -0,0 +1,172 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Preference-row atomicity and dynamic batching tests.""" + +import hashlib +import inspect +from argparse import Namespace +from pathlib import Path + +import pytest + + +try: + from relax.backends.megatron import data as data_module + from relax.backends.megatron.data import expand_preference_rollout_data + from relax.utils.training.preference_utils import pack_preference_pair_indices +except Exception as exc: + pytest.skip(f"relax.backends.megatron unavailable: {exc}", allow_module_level=True) + + +def _pair_rows(count=2): + return { + "pair_ids": list(range(100, 100 + count)), + "chosen_tokens": [[1, 2]] * count, + "rejected_tokens": [[1, 3]] * count, + "chosen_loss_masks": [[0, 1]] * count, + "rejected_loss_masks": [[0, 1]] * count, + "chosen_total_lengths": [2] * count, + "rejected_total_lengths": [2] * count, + } + + +def test_expand_keeps_pairs_atomic_and_preserves_dynamic_denominator(): + rows = _pair_rows() + rows["dynamic_global_batch_size"] = 2 + flat = expand_preference_rollout_data(rows) + assert flat["dynamic_global_batch_size"] == 2 + assert flat["preference_branch_pair_ids"] == [100, 100, 101, 101] + assert flat["preference_is_chosen"] == [True, False, True, False] + assert flat["preference_pair_costs"] == [4, 4] + + +def test_capacity_packer_is_deterministic_complete_and_bounded(): + costs = [2, 4, 4, 5, 5] + first = pack_preference_pair_indices(costs, ["a", "b", "c", "d", "e"], capacity=10) + second = pack_preference_pair_indices(costs, ["a", "b", "c", "d", "e"], capacity=10) + assert first == second + assert sorted(index for group in first for index in group) == list(range(len(costs))) + assert all(sum(costs[index] for index in group) <= 10 for group in first) + + +def test_preference_iterator_validates_step_global_pair_denominator(monkeypatch): + flat = expand_preference_rollout_data(_pair_rows()) + monkeypatch.setattr(data_module.mpu, "get_data_parallel_world_size", lambda **kwargs: 1) + monkeypatch.setattr(data_module.mpu, "get_data_parallel_group_gloo", lambda **kwargs: object()) + + def all_gather_object(output, value, **_kwargs): + output[:] = [value] + + monkeypatch.setattr(data_module.dist, "all_gather_object", all_gather_object) + args = Namespace(global_batch_size=2, max_tokens_per_gpu=16) + iterators, counts = data_module._get_preference_data_iterator(args, flat, None) + assert counts == [1] + assert flat["dynamic_global_batch_size"] == 2 + assert len(iterators) == 1 + + invalid = expand_preference_rollout_data(_pair_rows()) + invalid["dynamic_global_batch_size"] = 4 + with pytest.raises(ValueError, match="step-global preference pair count"): + data_module._get_preference_data_iterator(args, invalid, None) + + +def test_dp2_pair_rows_remain_atomic_with_global_pair_denominator(monkeypatch): + flat = expand_preference_rollout_data(_pair_rows()) + monkeypatch.setattr(data_module.mpu, "get_data_parallel_world_size", lambda **kwargs: 2) + monkeypatch.setattr(data_module.mpu, "get_data_parallel_group_gloo", lambda **kwargs: object()) + + def all_gather_object(output, value, **_kwargs): + output[:] = [value, value] + + monkeypatch.setattr(data_module.dist, "all_gather_object", all_gather_object) + args = Namespace(global_batch_size=4, max_tokens_per_gpu=4) + iterators, counts = data_module._get_preference_data_iterator(args, flat, None) + assert flat["dynamic_global_batch_size"] == 4 + assert counts == [2] + seen = [] + for _ in range(counts[0]): + batch = iterators[0].get_next(["preference_branch_pair_ids", "preference_is_chosen"]) + assert len(set(batch["preference_branch_pair_ids"])) == 1 + assert set(batch["preference_is_chosen"]) == {False, True} + seen.extend(batch["preference_branch_pair_ids"]) + assert sorted(seen) == [100, 100, 101, 101] + + +def test_preference_iterator_has_no_device_scalar_readback(): + source = inspect.getsource(data_module._get_preference_data_iterator) + assert ".item(" not in source + assert "all_reduce(" not in source + + +def test_preference_iterator_rejects_unequal_dp_pair_rows_via_gloo(monkeypatch): + flat = expand_preference_rollout_data(_pair_rows()) + monkeypatch.setattr(data_module.mpu, "get_data_parallel_world_size", lambda **kwargs: 2) + monkeypatch.setattr(data_module.mpu, "get_data_parallel_group_gloo", lambda **kwargs: object()) + + def all_gather_object(output, value, **_kwargs): + output[:] = [value, (value[0] + 1, value[1], value[2])] + + monkeypatch.setattr(data_module.dist, "all_gather_object", all_gather_object) + args = Namespace(global_batch_size=4, max_tokens_per_gpu=16) + with pytest.raises(ValueError, match="equal local pair rows"): + data_module._get_preference_data_iterator(args, flat, None) + + +def test_oversize_error_names_pair_cost_and_capacity(): + with pytest.raises(ValueError, match=r"oversize preference pair 'pair-x'.*cost 11, capacity=10"): + pack_preference_pair_indices([11], ["pair-x"], capacity=10) + + +def test_pinned_seqlen_sampler_consumes_pair_costs_and_keeps_equal_dp_groups(): + """Exercise the Docker-pinned TransferQueue sampler, not a local stand- + in.""" + transfer_queue = pytest.importorskip("transfer_queue") + sampler_type = transfer_queue.SeqlenBalancedSampler + source_path = Path(inspect.getsourcefile(sampler_type) or "") + normalized_source = source_path.read_bytes().replace(b"\r\n", b"\n") + assert hashlib.sha256(normalized_source).hexdigest() == ( + "dc6c2db50df4b9448d4845ccacef67a400517db892b5cd55de2e22f6baf6888b" + ) + + class PairPartition: + def __init__(self, pair_rows): + self.requested_indexes = None + self.metadata = { + index: {"total_lengths": len(row["chosen_tokens"]) + len(row["rejected_tokens"])} + for index, row in enumerate(pair_rows) + } + + def get_custom_meta(self, indexes): + self.requested_indexes = list(indexes) + return {index: self.metadata[index] for index in indexes} + + rows = [ + {"pair_id": 100, "chosen_tokens": list(range(60)), "rejected_tokens": list(range(40))}, + {"pair_id": 101, "chosen_tokens": list(range(50)), "rejected_tokens": list(range(40))}, + {"pair_id": 102, "chosen_tokens": list(range(6)), "rejected_tokens": list(range(4))}, + {"pair_id": 103, "chosen_tokens": [0], "rejected_tokens": [1]}, + ] + partition = PairPartition(rows) + sampler = sampler_type(n_samples_per_prompt=1, dp_size=2) + assignments = [] + for rank in range(2): + sampled, consumed = sampler.sample( + [0, 1, 2, 3], + batch_size=2, + task_name="dpo", + partition_id="train_0", + dp_rank=rank, + batch_index=0, + partition=partition, + ) + assert sampled == consumed + assert len(sampled) == 2 + assignments.append(sampled) + for index in sampled: + assert rows[index]["pair_id"] == 100 + index + assert set(rows[index]) == {"pair_id", "chosen_tokens", "rejected_tokens"} + + assert partition.requested_indexes == [0, 1, 2, 3] + assert sorted(index for rank_rows in assignments for index in rank_rows) == [0, 1, 2, 3] + rank_costs = [sum(partition.metadata[index]["total_lengths"] for index in rank_rows) for rank_rows in assignments] + assert sorted(rank_costs) == [100, 102] diff --git a/tests/backends/megatron/test_sft_train_data_fields.py b/tests/backends/megatron/test_sft_train_data_fields.py index 21a58d111..33dc5d1f0 100644 --- a/tests/backends/megatron/test_sft_train_data_fields.py +++ b/tests/backends/megatron/test_sft_train_data_fields.py @@ -5,6 +5,9 @@ from argparse import Namespace +import pytest +import torch + def _mk_actor_args(loss_type: str): return Namespace( @@ -41,6 +44,49 @@ def test_sft_data_fields_excludes_rl_only_keys(): assert forbidden not in fields, f"SFT data_fields leaked RL key: {forbidden}" +def test_preference_data_fields_keep_pairs_atomic(): + from relax.utils.training.data_fields import build_data_fields + + args = _mk_actor_args(loss_type="sft") + args.sft_objective = "dpo" + + fields = build_data_fields(args) + + assert fields == [ + "pair_ids", + "chosen_tokens", + "rejected_tokens", + "chosen_loss_masks", + "rejected_loss_masks", + "chosen_total_lengths", + "rejected_total_lengths", + ] + + +def test_preference_rows_expand_before_generic_rollout_post_processing(monkeypatch): + pytest.importorskip("megatron.core") + from relax.utils.data import stream_dataloader + + rollout_data = { + "pair_ids": [17], + "chosen_tokens": [[1, 2, 3]], + "rejected_tokens": [[1, 4]], + "chosen_loss_masks": [[0, 1, 1]], + "rejected_loss_masks": [[0, 1]], + "chosen_total_lengths": [3], + "rejected_total_lengths": [2], + } + args = Namespace(qkv_format="thd", is_vl_model=False, uses_unsplit_forward=False, use_opd=False) + monkeypatch.setattr(stream_dataloader.device_utils, "make_current_torch_device", lambda: torch.device("cpu")) + + stream_dataloader.post_process_rollout_data(args, rollout_data) + + assert [tensor.tolist() for tensor in rollout_data["tokens"]] == [[1, 2, 3], [1, 4]] + assert [tensor.tolist() for tensor in rollout_data["loss_masks"]] == [[0, 1, 1], [0, 1]] + assert rollout_data["preference_pair_ids"] == [17] + assert rollout_data["preference_pair_costs"] == [5] + + def test_rl_data_fields_unchanged(): """RL path must keep the existing field set.""" from relax.utils.training.data_fields import build_data_fields diff --git a/tests/data/test_prepare_ultrafeedback_preferences.py b/tests/data/test_prepare_ultrafeedback_preferences.py new file mode 100644 index 000000000..93ae94e40 --- /dev/null +++ b/tests/data/test_prepare_ultrafeedback_preferences.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Schema validation for the deterministic UltraFeedback preparation script.""" + +import pytest + +from scripts.data.prepare_ultrafeedback_preferences import _select, _validate_row + + +def _row() -> dict: + return { + "prompt_id": "pair-1", + "chosen": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "good"}, + ], + "rejected": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "bad"}, + ], + } + + +@pytest.mark.parametrize( + ("replacement", "match"), + [ + ("not-a-message", "message object"), + ({"content": "question"}, "invalid role"), + ({"role": "invalid", "content": "question"}, "invalid role"), + ({"role": "user", "content": 123}, "content must be a string"), + ], +) +@pytest.mark.parametrize("branch", ["chosen", "rejected"]) +def test_validate_row_rejects_invalid_message_schema(replacement, match: str, branch: str): + row = _row() + row[branch][0] = replacement + + with pytest.raises(ValueError, match=match): + _validate_row(row, split="train_prefs", index=0) + + +def test_select_classifies_non_object_rows_as_schema_rejections(): + selected, rejected = _select(["not-an-object", _row()], split="train_prefs", count=1) + + assert selected[0]["prompt_id"] == "pair-1" + assert rejected == [ + { + "prompt_id": None, + "source_index": 0, + "reason_code": "schema", + "reason": "train_prefs[0] must be an object", + } + ] diff --git a/tests/engine/sft/dataset/test_preference.py b/tests/engine/sft/dataset/test_preference.py new file mode 100644 index 000000000..d67aadbc0 --- /dev/null +++ b/tests/engine/sft/dataset/test_preference.py @@ -0,0 +1,223 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Preference-pair schema, rendering, truncation, and queue tests.""" + +import json +from pathlib import Path + +import pytest +import torch + +from relax.engine.sft.dataset.preference import ( + PreferenceDataError, + PreferenceStreamingDataset, + pack_preference_pairs_for_tq, +) + + +def _write_jsonl(path: Path, rows: list[dict]) -> None: + with path.open("w") as file: + for row in rows: + file.write(json.dumps(row) + "\n") + + +class _FakeTokenizer: + chat_template = "{% generation %}assistant{% endgeneration %}" + + def apply_chat_template( + self, + messages, + *, + tools=None, # noqa: ARG002 + tokenize=True, # noqa: ARG002 + return_tensors=None, # noqa: ARG002 + return_dict=False, + return_assistant_tokens_mask=False, + **kwargs, # noqa: ARG002 + ): + ids: list[int] = [] + masks: list[int] = [] + for message in messages: + prefix = {"system": 10, "user": 20, "assistant": 30}[message["role"]] + content = message["content"] + encoded = [prefix + (ord(char) % 10) for char in content] + ids.extend(encoded) + masks.extend([int(message["role"] == "assistant")] * len(encoded)) + input_ids = torch.tensor([ids], dtype=torch.long) + if return_assistant_tokens_mask: + return {"input_ids": input_ids, "assistant_masks": [masks]} + return input_ids + + +def _dataset(path: Path, **kwargs) -> PreferenceStreamingDataset: + return PreferenceStreamingDataset( + path=str(path), + tokenizer=_FakeTokenizer(), + prompt_key="prompt", + chosen_key="chosen", + rejected_key="rejected", + pair_id_key="prompt_id", + prefetch_max_cached=0, + **kwargs, + ) + + +def test_explicit_pair_builds_identical_prompt_and_completion_only_masks(tmp_path: Path): + path = tmp_path / "pairs.jsonl" + _write_jsonl( + path, + [ + { + "prompt_id": "pair-1", + "prompt": [{"role": "user", "content": "question"}], + "chosen": {"role": "assistant", "content": "good"}, + "rejected": {"role": "assistant", "content": "bad"}, + } + ], + ) + + dataset = _dataset(path, max_length=32, max_completion_length=8, pair_capacity=64) + dataset.shuffle(0) + pairs, crossed = dataset.get_batch(1) + + assert crossed is False + assert len(pairs) == 1 + pair = pairs[0] + chosen_prompt = pair.chosen_tokens[: pair.chosen_prompt_length] + rejected_prompt = pair.rejected_tokens[: pair.rejected_prompt_length] + assert torch.equal(chosen_prompt, rejected_prompt) + assert pair.chosen_loss_mask[: pair.chosen_prompt_length].sum().item() == 0 + assert pair.rejected_loss_mask[: pair.rejected_prompt_length].sum().item() == 0 + assert pair.chosen_loss_mask[pair.chosen_prompt_length :].all() + assert pair.rejected_loss_mask[pair.rejected_prompt_length :].all() + + +def test_implicit_ultrafeedback_pair_extracts_strict_common_prefix(tmp_path: Path): + path = tmp_path / "pairs.jsonl" + _write_jsonl( + path, + [ + { + "prompt_id": "pair-1", + "chosen": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "good"}, + ], + "rejected": [ + {"role": "user", "content": "question"}, + {"role": "assistant", "content": "bad"}, + ], + } + ], + ) + + dataset = _dataset(path, max_length=32, max_completion_length=8, pair_capacity=64) + pair = dataset.get_processed_pair(0) + + assert pair.pair_id == "pair-1" + assert pair.chosen_completion_length == 4 + assert pair.rejected_completion_length == 3 + + +def test_rejection_reason_is_classified_counted_and_still_fail_fast(tmp_path: Path): + path = tmp_path / "identical.jsonl" + _write_jsonl( + path, + [ + { + "prompt_id": "pair-identical", + "prompt": [{"role": "user", "content": "question"}], + "chosen": {"role": "assistant", "content": "same"}, + "rejected": {"role": "assistant", "content": "same"}, + } + ], + ) + dataset = _dataset(path) + with pytest.raises(PreferenceDataError) as exc_info: + dataset.get_processed_pair(0) + assert exc_info.value.reason_code == "identical" + assert dataset.rejection_counts == {"identical": 1} + assert dataset.rejection_records[0]["pair_id"] == "pair-identical" + + +@pytest.mark.parametrize( + ("update", "match"), + [ + ({"prompt_id": None}, "prompt_id"), + ({"rejected": {"role": "user", "content": "bad"}}, "assistant"), + ({"rejected": {"role": "assistant", "content": "good"}}, "identical"), + ], +) +def test_pair_schema_rejects_invalid_rows(tmp_path: Path, update: dict, match: str): + row = { + "prompt_id": "pair-1", + "prompt": [{"role": "user", "content": "question"}], + "chosen": {"role": "assistant", "content": "good"}, + "rejected": {"role": "assistant", "content": "bad"}, + } + row.update(update) + path = tmp_path / "pairs.jsonl" + _write_jsonl(path, [row]) + + with pytest.raises(ValueError, match=match): + _dataset(path, max_length=32, max_completion_length=8, pair_capacity=64).get_processed_pair(0) + + +def test_preference_dataset_rejects_duplicate_pair_ids(tmp_path: Path): + row = { + "prompt_id": "duplicate", + "prompt": [{"role": "user", "content": "question"}], + "chosen": {"role": "assistant", "content": "good"}, + "rejected": {"role": "assistant", "content": "bad"}, + } + path = tmp_path / "pairs.jsonl" + _write_jsonl(path, [row, row]) + + with pytest.raises(ValueError, match="duplicate preference pair ID"): + _dataset(path, max_length=32, max_completion_length=8, pair_capacity=64) + + +def test_shared_prompt_and_completion_truncation_preserves_pair_difference(tmp_path: Path): + path = tmp_path / "pairs.jsonl" + _write_jsonl( + path, + [ + { + "prompt_id": "pair-1", + "prompt": [{"role": "user", "content": "0123456789"}], + "chosen": {"role": "assistant", "content": "chosen"}, + "rejected": {"role": "assistant", "content": "reject"}, + } + ], + ) + + pair = _dataset(path, max_length=8, max_completion_length=3, pair_capacity=16).get_processed_pair(0) + + assert pair.chosen_prompt_length == pair.rejected_prompt_length == 5 + assert pair.chosen_completion_length == pair.rejected_completion_length == 3 + assert pair.chosen_total_length + pair.rejected_total_length == 16 + assert not torch.equal( + pair.chosen_tokens[pair.chosen_prompt_length :], pair.rejected_tokens[pair.rejected_prompt_length :] + ) + + +def test_pack_pair_rows_and_custom_meta_are_aligned(tmp_path: Path): + path = tmp_path / "pairs.jsonl" + rows = [] + for idx in range(2): + rows.append( + { + "prompt_id": f"pair-{idx}", + "prompt": [{"role": "user", "content": f"q{idx}"}], + "chosen": {"role": "assistant", "content": f"yes{idx}"}, + "rejected": {"role": "assistant", "content": f"no{idx}"}, + } + ) + _write_jsonl(path, rows) + dataset = _dataset(path, max_length=16, max_completion_length=8, pair_capacity=32) + + batch, custom_meta = pack_preference_pairs_for_tq([dataset.get_processed_pair(0), dataset.get_processed_pair(1)]) + + assert len(batch["pair_ids"]) == len(custom_meta) == 2 + for idx, metadata in enumerate(custom_meta): + assert metadata["total_lengths"] == (batch["chosen_total_lengths"][idx] + batch["rejected_total_lengths"][idx]) diff --git a/tests/engine/sft/test_preference_runtime.py b/tests/engine/sft/test_preference_runtime.py new file mode 100644 index 000000000..80d8f1941 --- /dev/null +++ b/tests/engine/sft/test_preference_runtime.py @@ -0,0 +1,93 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Fail-fast validation for offline preference objectives.""" + +from argparse import Namespace + +import pytest + +from relax.engine.sft.runtime import is_preference_mode, validate_preference_args + + +def _args(**overrides) -> Namespace: + values = { + "loss_type": "sft", + "sft_objective": "dpo", + "custom_dataset_class_path": None, + "multimodal_keys": None, + "n_samples_per_prompt": 1, + "tensor_model_parallel_size": 1, + "pipeline_model_parallel_size": 1, + "context_parallel_size": 1, + "dynamic_context_parallel": False, + "qkv_format": "thd", + "fully_async": False, + "hybrid": False, + "use_gloo_process_groups": True, + "sft_chunked_logits": False, + "enable_mtp_training": False, + "calculate_per_token_loss": False, + "lora_rank": 0, + "hidden_dropout": 0.0, + "attention_dropout": 0.0, + "sft_predict_interval": None, + "eval_interval": None, + "eval_prompt_data": None, + "eval_size": None, + "dpo_beta": 0.1, + "rollout_temperature": 1.0, + "dpo_reference_free": False, + "dpo_reference_repository": "Qwen/Qwen3-0.6B", + "dpo_reference_revision": "fixed-revision", + "ref_load": None, + "ref_update_interval": None, + "enable_weights_backuper": True, + "preference_max_length": 1024, + "preference_max_completion_length": 512, + "seq_length": 2048, + } + values.update(overrides) + return Namespace(**values) + + +def test_preference_mode_is_nested_under_sft(): + assert is_preference_mode(_args()) + assert not is_preference_mode(_args(loss_type="policy_loss")) + assert not is_preference_mode(_args(sft_objective="causal_lm")) + assert not is_preference_mode(_args(sft_objective="reward_model")) + + +@pytest.mark.parametrize( + ("overrides", "match"), + [ + ({"n_samples_per_prompt": 2}, "n-samples-per-prompt"), + ({"tensor_model_parallel_size": 2}, "TP=CP=PP=1"), + ({"context_parallel_size": 2}, "TP=CP=PP=1"), + ({"dynamic_context_parallel": True}, "dynamic context"), + ({"qkv_format": "bshd"}, "qkv-format thd"), + ({"use_gloo_process_groups": False}, "use-gloo-process-groups"), + ({"lora_rank": 8}, "LoRA"), + ({"hidden_dropout": 0.1}, "dropout"), + ({"ref_update_interval": 10}, "frozen reference"), + ({"ref_load": "/tmp/ref"}, "do not use --ref-load"), + ({"dpo_reference_free": True, "ref_load": "/tmp/ref"}, "do not use --ref-load"), + ({"dpo_beta": float("nan")}, "finite and positive"), + ({"rollout_temperature": 0.8}, "rollout-temperature 1.0"), + ({"rollout_temperature": float("nan")}, "rollout-temperature 1.0"), + ({"rollout_temperature": float("inf")}, "rollout-temperature 1.0"), + ({"preference_max_completion_length": 2048}, "must not exceed"), + ({"eval_prompt_data": ["heldout", "eval.jsonl"]}, "follow-up reward-modeling PR"), + ], +) +def test_preference_validation_rejects_unsupported_configs(overrides: dict, match: str): + with pytest.raises(ValueError, match=match): + validate_preference_args(_args(**overrides)) + + +def test_reference_free_dpo_does_not_require_ref_update_constraint(): + validate_preference_args(_args(dpo_reference_free=True, ref_update_interval=10)) + + +def test_standard_dpo_requires_explicit_reference_repository_and_revision(): + with pytest.raises(ValueError, match="dpo-reference-repository"): + validate_preference_args(_args(dpo_reference_repository=None)) diff --git a/tests/utils/training/test_preference_utils.py b/tests/utils/training/test_preference_utils.py new file mode 100644 index 000000000..f1cf2d4a1 --- /dev/null +++ b/tests/utils/training/test_preference_utils.py @@ -0,0 +1,99 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Pure numerical and batching tests for offline preference training.""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from relax.utils.training.preference_utils import ( + build_causal_lm_labels, + dpo_pair_loss, + pack_preference_pair_indices, + require_tensor_condition, +) + + +def test_tensor_condition_uses_async_assert_without_python_bool_on_cuda(monkeypatch): + class _CudaCondition: + device = SimpleNamespace(type="cuda") + + def __bool__(self): + raise AssertionError("CUDA conditions must not be converted to Python bool") + + calls = [] + monkeypatch.setattr(torch, "_assert_async", lambda condition, message: calls.append((condition, message))) + condition = _CudaCondition() + + require_tensor_condition(condition, "finite") + + assert calls == [(condition, "finite")] + + +def test_build_causal_lm_labels_uses_next_token_mask(): + tokens = torch.tensor([10, 11, 12, 13, 14]) + raw_mask = torch.tensor([0, 0, 1, 1, 1]) + + labels = build_causal_lm_labels(tokens, raw_mask) + + assert labels.tolist() == [-100, 12, 13, 14, -100] + + +@pytest.mark.parametrize("beta", [0.01, 0.1, 1.0]) +def test_dpo_pair_loss_matches_independent_reference_and_gradient(beta: float): + policy_chosen = torch.tensor([-2.0, -0.25, 4.0], dtype=torch.float32, requires_grad=True) + policy_rejected = torch.tensor([-3.0, 0.75, -1.0], dtype=torch.float32, requires_grad=True) + ref_chosen = torch.tensor([-2.5, 0.5, 1.0], dtype=torch.float32) + ref_rejected = torch.tensor([-2.0, -0.5, -2.0], dtype=torch.float32) + + actual = dpo_pair_loss( + policy_chosen, + policy_rejected, + reference_chosen=ref_chosen, + reference_rejected=ref_rejected, + beta=beta, + ) + expected = -F.logsigmoid(beta * ((policy_chosen - policy_rejected) - (ref_chosen - ref_rejected))) + assert torch.allclose(actual, expected, rtol=1e-6, atol=1e-6) + + actual.sum().backward() + actual_grad = policy_chosen.grad.detach().clone() + policy_chosen.grad = None + expected.sum().backward() + assert torch.allclose(actual_grad, policy_chosen.grad, rtol=1e-6, atol=1e-6) + + +def test_dpo_pair_loss_reference_free_matches_independent_reference(): + chosen = torch.tensor([-1.0, 2.0], requires_grad=True) + rejected = torch.tensor([0.5, -2.0], requires_grad=True) + + actual = dpo_pair_loss(chosen, rejected, beta=0.1, reference_free=True) + expected = -F.logsigmoid(0.1 * (chosen - rejected)) + + assert torch.allclose(actual, expected, rtol=1e-6, atol=1e-6) + + +def test_dpo_pair_loss_rejects_missing_reference_and_non_finite_values(): + finite = torch.tensor([0.0]) + with pytest.raises(ValueError, match="reference log-probabilities"): + dpo_pair_loss(finite, finite) + with pytest.raises(ValueError, match="finite"): + dpo_pair_loss(torch.tensor([float("nan")]), finite, reference_free=True) + + +def test_pair_packer_is_deterministic_complete_and_capacity_safe(): + costs = [2, 4, 4, 5, 5] + pair_ids = ["a", "b", "c", "d", "e"] + + bins = pack_preference_pair_indices(costs, pair_ids, capacity=10) + + assert sorted(index for group in bins for index in group) == list(range(len(costs))) + assert all(sum(costs[index] for index in group) <= 10 for group in bins) + assert bins == pack_preference_pair_indices(costs, pair_ids, capacity=10) + + +def test_pair_packer_reports_oversize_pair(): + with pytest.raises(ValueError, match="oversize.*pair-a.*11"): + pack_preference_pair_indices([11], ["pair-a"], capacity=10)