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..a29f8006c --- /dev/null +++ b/docs/en/guide/dpo-training.md @@ -0,0 +1,72 @@ +# 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. + +## Reward modeling and acceptance artifacts + +The companion recipe is `scripts/training/reward_modeling/run-qwen3-0.6B-ultrafeedback-1xgpu.sh`. It defaults to 200 optimizer steps and 32 pairs per global batch. Preference evaluation runs before the first optimizer step (step 0), periodically, and after the final completed step even when the interval does not divide the run length. + +Both DPO and reward modeling write acceptance data under `//preference_eval/`: the canonical probe contract and SHA-256, the DP/micro-batch plan and SHA-256, step-0/final per-pair JSONL, and a 10,000-replicate FP64 PCG64 paired-bootstrap summary. Final evaluation fails if probe preprocessing, pair order, or the batch plan differs from step 0. Retain this directory together with the expanded command, environment inventory, raw stdout/stderr, metrics, and curves. + +Reward-model Megatron checkpoints persist `sft_objective=reward_model`, `head_type=reward_model_terminal_v1`, and `checkpoint_role=actor`. Resume rejects missing or incompatible metadata, non-exact scalar-head keys/shapes, partial optimizer/RNG restoration, and PPO critic checkpoints. diff --git a/docs/zh/guide/dpo-training.md b/docs/zh/guide/dpo-training.md new file mode 100644 index 000000000..0978c2095 --- /dev/null +++ b/docs/zh/guide/dpo-training.md @@ -0,0 +1,72 @@ +# 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。 + +## Reward Modeling 与验收产物 + +配套 recipe 为 `scripts/training/reward_modeling/run-qwen3-0.6B-ultrafeedback-1xgpu.sh`，默认运行 200 个 optimizer step，global batch 为 32 pairs。偏好评测会在第一次 optimizer step 之前（step 0）、周期边界以及最终完成 step 后运行；最终评测不依赖 interval 恰好整除训练步数。 + +DPO 与 Reward Modeling 都会在 `//preference_eval/` 下写出验收数据：canonical probe 合同及 SHA-256、DP/micro-batch plan 及 SHA-256、step-0/final 逐 pair JSONL，以及 10,000 次 FP64 PCG64 paired-bootstrap summary。若 final 与 step 0 的预处理、pair 顺序或 batch plan 不一致，评测会立即失败。提交证据时需将该目录与展开后的命令、环境清单、原始 stdout/stderr、metrics 和曲线一并保留。 + +Reward Model 的 Megatron checkpoint 会持久化 `sft_objective=reward_model`、`head_type=reward_model_terminal_v1` 与 `checkpoint_role=actor`。resume 会拒绝缺失或不兼容的 metadata、非精确 scalar-head key/shape、只恢复部分 optimizer/RNG 状态，以及 PPO critic checkpoint。 diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 75143efa6..5c91a0ef1 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,8 @@ 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 ( + evaluation_step_for_rollout, + is_preference_mode, is_sft_mode, sft_partition_id, sft_task_name, @@ -77,7 +80,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 +88,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 +97,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 +241,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 +281,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 +444,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( @@ -560,10 +765,10 @@ def _run_step_evaluation(self, rollout_id: int, *, end_update_weight: bool = Fal should_run_predict = has_rollout and should_run_sft_predict(self.args, rollout_id) try: if should_run_eval: - if dist.get_rank() == 0: + if rollout_id > 0 and dist.get_rank() == 0: run( self.data_system_client.async_clear_partition( - partition_id=sft_partition_id(self.args, rollout_id) + partition_id=sft_partition_id(self.args, rollout_id - 1) ) ) dist.barrier(group=get_gloo_group()) @@ -610,6 +815,14 @@ def _request_rollout_evaluation(self, rollout_id: int, *, end_update_weight: boo self._run_step_evaluation(rollout_id, end_update_weight=end_update_weight) def train(self, rollout_id: int) -> None: + if ( + rollout_id == 0 + and is_preference_mode(self.args) + and not self.args.debug_train_only + and should_run_sft_eval(self.args, 0) + ): + self._run_step_evaluation(0) + if self.args.offload_rollout and dist.get_rank() == 0: pre_train_offload_handles = [] if self.genrm_manager is not None: @@ -783,6 +996,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 +1015,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 +1033,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 +1048,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 +1077,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" @@ -974,7 +1200,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: # RL-only generative eval (uses SGLang via rollout_manager.eval). SFT # uses local eval/predict runner below. dist.barrier(group=get_gloo_group()) - self._run_step_evaluation(rollout_id) + self._run_step_evaluation(evaluation_step_for_rollout(self.args, rollout_id)) # On the final training step the rollout component has already exited # its main loop, so nothing else awaits the eval handler. Block here @@ -1443,7 +1669,10 @@ def train_hybrid(self, rollout_id) -> None: self.update_weights() tracking_utils.flush_metrics(self.args, compute_rollout_step(self.args, rollout_id)) dist.barrier(group=get_gloo_group()) - self._run_step_evaluation(rollout_id, end_update_weight=True) + self._run_step_evaluation( + evaluation_step_for_rollout(self.args, rollout_id), + end_update_weight=True, + ) # On the final training step the rollout component has already exited # its main loop, so the eval just triggered above will not be awaited @@ -1526,7 +1755,10 @@ def train_async(self, rollout_id) -> None: rollout_only, actor_fwd_only = self._check_services_health() self.update_weights_fully_async(rollout_id, rollout_only=rollout_only, actor_fwd_only=actor_fwd_only) dist.barrier(group=get_gloo_group()) - self._run_step_evaluation(rollout_id, end_update_weight=True) + self._run_step_evaluation( + evaluation_step_for_rollout(self.args, rollout_id), + end_update_weight=True, + ) # On the final training step the rollout component has already # exited its main loop, so the eval just triggered above will not # be awaited anywhere. Block until it finishes; otherwise the @@ -1596,9 +1828,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 +2101,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..c9cfadd46 100644 --- a/relax/backends/megatron/checkpoint.py +++ b/relax/backends/megatron/checkpoint.py @@ -99,7 +99,121 @@ def _init_from_local_shards_and_global_metadata( # type: ignore[override] logger = get_logger(__name__) -__all__ = ["save_checkpoint"] +REWARD_MODEL_HEAD_TYPE = "reward_model_terminal_v1" + +__all__ = ["REWARD_MODEL_HEAD_TYPE", "load_checkpoint", "save_checkpoint", "scheduler_state_was_restored"] + + +def scheduler_state_was_restored(args, resumed_from_megatron: bool) -> bool: + return bool( + resumed_from_megatron + and not getattr(args, "no_load_optim", False) + and not getattr(args, "finetune", False) + and not getattr(args, "reset_optimizer_states", False) + ) + + +def _checkpoint_iteration_dir(load_path: str | Path, ckpt_step: int | None = None) -> Path: + path = Path(load_path) + if re.fullmatch(r"iter_\d{7}", path.name): + return path + tracker = path / "latest_checkpointed_iteration.txt" + try: + metadata = tracker.read_text(encoding="utf-8").strip() + except OSError as exc: + raise RuntimeError(f"cannot resolve Megatron checkpoint iteration from {tracker}") from exc + if metadata == "release": + return path / "release" + try: + iteration = int(metadata) + except ValueError as exc: + raise RuntimeError(f"cannot resolve Megatron checkpoint iteration from {tracker}") from exc + if ckpt_step is not None: + iteration = int(ckpt_step) + if iteration < 0: + raise RuntimeError(f"Megatron checkpoint iteration must be non-negative, got {iteration}") + return path / f"iter_{iteration:07d}" + + +def _metadata_value(metadata, name: str): + return metadata.get(name) if isinstance(metadata, dict) else getattr(metadata, name, None) + + +def _validate_reward_model_tensor_metadata(tensor_metadata: dict, hidden_size: int) -> None: + head_keys = [ + str(key) + for key in tensor_metadata + if not str(key).startswith("optimizer.") and ("output_layer." in str(key) or "reward_model_head." in str(key)) + ] + weight_keys = [key for key in head_keys if key.endswith("output_layer.weight")] + if len(weight_keys) != 1: + raise RuntimeError( + f"RM resume requires exactly one output_layer.weight checkpoint tensor, found {weight_keys}" + ) + unexpected = [key for key in head_keys if key != weight_keys[0]] + if unexpected: + raise RuntimeError(f"RM checkpoint contains unexpected scalar-head tensors: {unexpected}") + entry = tensor_metadata[weight_keys[0]] + shape = tuple(getattr(entry, "global_shape", getattr(entry, "shape", ()))) + expected = (1, int(hidden_size)) + if shape != expected: + raise RuntimeError(f"RM output_layer.weight shape mismatch: checkpoint={shape}, expected={expected}") + + +def _validate_checkpoint_contract(args, ddp_model, checkpoint_dir: Path) -> None: + from megatron.core import dist_checkpointing + + role = getattr(ddp_model[0], "role", "actor") + current_is_rm = ( + role == "actor" + and getattr(args, "loss_type", None) == "sft" + and getattr(args, "sft_objective", "causal_lm") == "reward_model" + ) + if current_is_rm and checkpoint_dir.name == "release": + raise RuntimeError( + "RM resume rejects release checkpoints because optimizer, scheduler, and RNG state are absent" + ) + if not dist_checkpointing.check_is_distributed_checkpoint(str(checkpoint_dir)): + if current_is_rm: + raise RuntimeError( + "RM resume requires a distributed checkpoint with contract metadata; legacy Megatron " + "checkpoints are not supported" + ) + return + common = dist_checkpointing.load_common_state_dict(checkpoint_dir) + saved_args = common.get("args") + if saved_args is None: + raise RuntimeError(f"checkpoint {checkpoint_dir} is missing saved args metadata") + saved_objective = _metadata_value(saved_args, "sft_objective") + saved_head_type = _metadata_value(saved_args, "head_type") + saved_role = _metadata_value(saved_args, "checkpoint_role") + saved_is_rm = saved_objective == "reward_model" or saved_head_type == REWARD_MODEL_HEAD_TYPE + + if current_is_rm: + if (saved_objective, saved_head_type, saved_role) != ( + "reward_model", + REWARD_MODEL_HEAD_TYPE, + "actor", + ): + raise RuntimeError( + "RM resume requires checkpoint metadata " + "sft_objective=reward_model, head_type=reward_model_terminal_v1, checkpoint_role=actor; " + f"got objective={saved_objective!r}, head_type={saved_head_type!r}, role={saved_role!r}" + ) + incompatible_flags = [ + name + for name in ("no_load_optim", "no_load_rng", "finetune", "reset_optimizer_states") + if bool(getattr(args, name, False)) + ] + if incompatible_flags: + raise RuntimeError( + f"RM resume must restore optimizer, scheduler, and RNG state; incompatible flags: {incompatible_flags}" + ) + tensor_metadata = dist_checkpointing.load_tensors_metadata(str(checkpoint_dir)) + _validate_reward_model_tensor_metadata(tensor_metadata, int(args.hidden_size)) + elif saved_is_rm: + target = "PPO critic" if role == "critic" else "non-RM actor/SFT" + raise RuntimeError(f"{target} load rejects reward-model checkpoints") def _alias_renamed_transfer_queue_enum() -> None: @@ -130,8 +244,13 @@ 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() + _validate_checkpoint_contract( + args, + ddp_model, + _checkpoint_iteration_dir(load_path, getattr(args, "ckpt_step", None)), + ) try: return _load_checkpoint_megatron( ddp_model=ddp_model, @@ -190,7 +309,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..fa312db45 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 ( @@ -449,6 +451,7 @@ def get_batch( from relax.utils.sft_utils import align_loss_mask_for_sft + batch["raw_loss_masks"] = list(batch["loss_masks"]) loss_masks: list[torch.Tensor] = [] per_sample_loss_masks: list[torch.Tensor] = [] full_per_sample_loss_masks: list[torch.Tensor] = [] @@ -703,6 +706,146 @@ 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", + "chosen_score_positions", + "rejected_score_positions", + ) + 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": [], + "score_positions": [], + "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]) + score_position = int(rollout_data[f"{prefix}_score_positions"][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" + ) + if not 0 <= score_position < total_length: + raise ValueError(f"preference pair row {index} {prefix} score position is out of range") + flat["tokens"].append(tokens) + flat["loss_masks"].append(loss_mask) + flat["total_lengths"].append(total_length) + flat["response_lengths"].append(total_length) + flat["score_positions"].append(score_position) + 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 +865,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 +1063,11 @@ 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", + "score_positions", ]: 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..69aba5333 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,13 @@ 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, + reward_model_pair_loss, + select_packed_sequence_scores, +) from relax.utils.types import RolloutBatch from .cp_utils import ( @@ -1265,6 +1273,155 @@ 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 reward_model_loss_function( + args: Namespace, # noqa: ARG001 + batch: RolloutBatch, + logits: torch.Tensor, + sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor], # noqa: ARG001 +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Compute pair-summed Bradley-Terry loss over adjacent chosen/rejected + branches.""" + if len(batch["total_lengths"]) % 2 != 0: + raise ValueError("reward-model micro-batch must contain an even number of chosen/rejected branches") + score_positions = batch.get("score_positions") + if score_positions is None: + raise ValueError("reward-model batch is missing score_positions") + scores = select_packed_sequence_scores( + logits, + batch["total_lengths"], + score_positions, + raw_loss_masks=batch.get("raw_loss_masks"), + packed_tokens=batch.get("tokens"), + branch_tokens=batch.get("unconcat_tokens"), + cu_seqlens=(batch["packed_seq_params"].cu_seqlens_q if batch.get("packed_seq_params") is not None else None), + ) + 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("reward-model 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) + chosen_scores = scores.index_select(0, chosen_index) + rejected_scores = scores.index_select(0, rejected_index) + pair_losses = reward_model_pair_loss(chosen_scores, rejected_scores) + margins = chosen_scores - rejected_scores + loss = pair_losses.sum() + if pair_losses.numel() == 0: + loss = loss + 0 * logits.sum() + return loss, { + "rm/loss": pair_losses.detach().sum(), + "rm/score_chosen_mean": chosen_scores.detach().sum(), + "rm/score_rejected_mean": rejected_scores.detach().sum(), + "rm/score_margin_mean": margins.detach().sum(), + "rm/accuracy": (margins > 0).to(torch.float32).detach().sum(), + "rm/_score_chosen_second_moment": chosen_scores.detach().square().sum(), + "rm/_score_rejected_second_moment": rejected_scores.detach().square().sum(), + } + + def sft_loss_function_chunked( args: Namespace, batch: RolloutBatch, @@ -1364,6 +1521,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 +1544,11 @@ 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_objective", "causal_lm") == "reward_model": + func = reward_model_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..d3ad5449d 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -27,7 +27,7 @@ from megatron.training.training import get_model from relax.backends.megatron.checkpoint import _save_lora_to_checkpoint -from relax.engine.sft.runtime import is_sft_mode +from relax.engine.sft.runtime import is_preference_mode, is_sft_mode from relax.utils import tracking_utils from relax.utils.data.stream_dataloader import StreamingTQIterator from relax.utils.env import Envs @@ -42,9 +42,16 @@ maybe_verify_critic_value_head_movement, release_critic_lm_heads, validate_critic_value_head_registration, + validate_reward_model_head_registration, ) -from .checkpoint import load_checkpoint, save_checkpoint +from .checkpoint import ( + REWARD_MODEL_HEAD_TYPE, + is_megatron_checkpoint, + load_checkpoint, + save_checkpoint, + scheduler_state_was_restored, +) from .data import DataIterator, get_batch from .loss import loss_function from .model_provider import get_model_provider_func, wrap_model_provider_with_freeze @@ -696,6 +703,19 @@ def force_param_sync(model_chunks: Sequence[DDP]) -> None: model_chunk.start_param_sync(force_sync=True) +def _restore_micro_batch_output_order(values: list, micro_batch_indices: list[list[int]]) -> list: + """Restore per-sample outputs from packed micro-batch order.""" + origin_indices = sum(micro_batch_indices, []) + if len(values) != len(origin_indices): + return values + if sorted(origin_indices) != list(range(len(origin_indices))): + raise RuntimeError("micro-batch indices must be a complete permutation of original sample indices") + origin_values = [None] * len(values) + for value, origin_index in zip(values, origin_indices, strict=False): + origin_values[origin_index] = value + return origin_values + + @torch.no_grad() def forward_only( f: Callable[..., dict[str, list[torch.Tensor]]], @@ -766,6 +786,7 @@ def forward_step( "multimodal_train_inputs", "total_lengths", "response_lengths", + "score_positions", "max_seq_lens", ], args.data_pad_size_multiplier, @@ -846,6 +867,13 @@ def forward_step( max_seq_lens=batch.get("max_seq_lens", None), padded_total_lengths=batch.get("padded_total_lengths", None), loss_masks=batch.get("loss_masks", None), + raw_loss_masks=batch.get("raw_loss_masks", None), + packed_tokens=batch.get("tokens", None), + branch_tokens=batch.get("unconcat_tokens", None), + cu_seqlens=( + batch["packed_seq_params"].cu_seqlens_q if batch.get("packed_seq_params") is not None else None + ), + score_positions=batch.get("score_positions", None), dynamic_cp_size=batch.get("dynamic_cp_size", None), dynamic_cp_rank=batch.get("dynamic_cp_rank", None), ) @@ -924,21 +952,17 @@ def _inject_meta(logits, *, _orig=f_partial, _meta=dcp_meta): assert isinstance(value[key], list) values += value[key] - if args.use_dynamic_batch_size and per_sample_output: + micro_batch_indices = data_iterator[0].micro_batch_indices + if micro_batch_indices is not None: # TODO: This is ugly... Find a better way to make the data have the same order. # TODO: move this out of the loop. - origin_indices = sum(data_iterator[0].micro_batch_indices, []) # Per-sample callbacks (log_probs/values) emit one tensor per # sample, so values aligns with origin_indices and we can # restore the pre-balance order. Per-microbatch callbacks # (e.g. compute_sft_eval_step) emit one aggregate per # microbatch — len(values) == num_microbatches, not # num_samples — and have no per-sample order to restore. - if len(values) == len(origin_indices): - origin_values = [None] * len(values) - for value, origin_index in zip(values, origin_indices, strict=False): - origin_values[origin_index] = value - values = origin_values + values = _restore_micro_batch_output_order(values, micro_batch_indices) rollout_data[f"{store_prefix}{key}"] = values return rollout_data @@ -1024,9 +1048,12 @@ def forward_step( "loss_masks", "log_probs", "ref_log_probs", + "preference_branch_pair_ids", + "preference_is_chosen", "values", "advantages", "returns", + "score_positions", "rollout_log_probs", "max_seq_lens", *_opd_keys, @@ -1256,6 +1283,13 @@ def forward_step( # CP degree under dynamic CP (and is a no-op under static CP, where the # count previously carried the cancelling cp factor). loss_reduced[key] = value / num_samples_or_tokens + if "rm/_score_chosen_second_moment" in loss_reduced: + chosen_second = loss_reduced.pop("rm/_score_chosen_second_moment") + rejected_second = loss_reduced.pop("rm/_score_rejected_second_moment") + chosen_mean = loss_reduced["rm/score_chosen_mean"] + rejected_mean = loss_reduced["rm/score_rejected_mean"] + loss_reduced["rm/score_chosen_std"] = math.sqrt(max(chosen_second - chosen_mean**2, 0.0)) + loss_reduced["rm/score_rejected_std"] = math.sqrt(max(rejected_second - rejected_mean**2, 0.0)) return loss_reduced, grad_norm return {}, grad_norm @@ -1515,6 +1549,14 @@ def save( opt_param_scheduler (OptimizerParamScheduler): LR/WD scheduler. """ args = get_args() + role = getattr(model[0], "role", "actor") + args.checkpoint_role = role + if role == "actor" and is_preference_mode(args) and args.sft_objective == "reward_model": + args.head_type = REWARD_MODEL_HEAD_TYPE + elif role == "critic": + args.head_type = "critic_value_terminal_v1" + else: + args.head_type = "causal_lm_v1" if should_disable_forward_pre_hook(args): disable_forward_pre_hook(model) save_checkpoint( @@ -1755,9 +1797,13 @@ def initialize_model_and_optimizer( model, optimizer, opt_param_scheduler = setup_model_and_optimizer(args, role) model[0].role = role value_head_param_ids = () + reward_head_param_ids = () if role == "critic": value_head_param_ids = validate_critic_value_head_registration(model, optimizer) + elif is_preference_mode(args) and args.sft_objective == "reward_model": + reward_head_param_ids = validate_reward_model_head_registration(model, optimizer) clear_memory() + resumed_from_megatron = is_megatron_checkpoint(args.load) iteration, _ = load_checkpoint( model, optimizer, @@ -1772,6 +1818,15 @@ def initialize_model_and_optimizer( "critic value head parameter identities changed during checkpoint loading" ) install_critic_value_head_runtime_check(model) + elif is_preference_mode(args) and args.sft_objective == "reward_model": + release_critic_lm_heads(model) + loaded_reward_head_param_ids = validate_reward_model_head_registration(model, optimizer) + assert loaded_reward_head_param_ids == reward_head_param_ids, ( + "reward-model head parameter identities changed during checkpoint loading" + ) clear_memory() + scheduler_was_restored = scheduler_state_was_restored(args, resumed_from_megatron) + if opt_param_scheduler is not None and not scheduler_was_restored: + opt_param_scheduler.step(increment=iteration * args.global_batch_size) return model, optimizer, opt_param_scheduler, iteration diff --git a/relax/backends/megatron/model_provider.py b/relax/backends/megatron/model_provider.py index 7e9805608..a5d031786 100644 --- a/relax/backends/megatron/model_provider.py +++ b/relax/backends/megatron/model_provider.py @@ -27,7 +27,10 @@ from relax.utils.logging_utils import get_logger from relax.utils.megatron_peft_utils import build_lora_peft, count_adapter_parameters, is_lora_enabled from relax.utils.misc import load_function -from relax.utils.training.ppo_utils import install_critic_value_head_in_provider +from relax.utils.training.ppo_utils import ( + install_critic_value_head_in_provider, + install_reward_model_head_in_provider, +) from .conditional_branch_sync import install_conditional_branch_sync @@ -197,6 +200,7 @@ def wrapped_model_provider( model = custom_model_provider(pre_process=pre_process, post_process=post_process) # Apply critic output layer if needed install_critic_value_head_in_provider(model, role, post_process) + install_reward_model_head_in_provider(model, args, role, post_process) _maybe_mark_unsplit_forward(args, model) install_conditional_branch_sync(args, model) _install_cp_probe(model) @@ -327,6 +331,7 @@ def provide_with_cp_probe(*p_args, **p_kwargs): model = original_provide(*p_args, **p_kwargs) post_process = p_kwargs.get("post_process", p_args[1] if len(p_args) > 1 else True) install_critic_value_head_in_provider(model, role, post_process, stash_lm_head=True) + install_reward_model_head_in_provider(model, args, role, post_process, stash_lm_head=True) _maybe_mark_unsplit_forward(args, model) install_conditional_branch_sync(args, model) _install_cp_probe(model) @@ -440,6 +445,7 @@ def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage model = GPTModel(**kwargs) install_critic_value_head_in_provider(model, role, post_process) + install_reward_model_head_in_provider(model, args, role, post_process) _maybe_mark_unsplit_forward(args, model) install_conditional_branch_sync(args, model) 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..9d868ffff 100644 --- a/relax/components/actor.py +++ b/relax/components/actor.py @@ -14,7 +14,13 @@ 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 ( + actor_training_input_ready, + 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 +84,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, ) ) @@ -249,9 +259,8 @@ def _wait_for_rollout_data(self) -> bool: True if data is ready and training can proceed, False if should continue waiting (caller should skip this iteration) """ - partition_id = sft_partition_id(self.config, self.step) partition_list = run(self.data_system_client.async_get_partition_list()) - if partition_list is None or partition_id not in partition_list: + if not actor_training_input_ready(self.config, self.step, partition_list): time.sleep(1) return False diff --git a/relax/components/sft.py b/relax/components/sft.py index 15ca2b15c..d33bb64ef 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, should_run_sft_eval 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, @@ -189,27 +216,47 @@ def _init_data_pipeline(self) -> None: eval_tool_key = getattr(self.config, "eval_tool_key", None) or self.config.tool_key # Eval is small + runs every `eval_interval`; disable prefetch so # we don't consume worker threads idly between eval rounds. - self._eval_dataset = SFTStreamingDataset( - path=[d.path for d in eval_prompt_data], - tokenizer=self._tokenizer, - processor_pool=self._processor_pool, - capacity=capacity, - prompt_key=eval_input_key, - label_key=eval_label_key, - multimodal_keys=self.config.multimodal_keys, - conversation_key_map=getattr(self.config, "conversation_key_map", None), - metadata_key=self.config.metadata_key, - tool_key=eval_tool_key, - system_prompt=self.config.system_prompt, - source_name="+".join(d.name for d in eval_prompt_data), - seed=seed, - prefetch_max_cached=0, - pad_token_ids=pad_token_ids, - oversize_strategy=oversize_strategy, - oversize_custom_fn=oversize_custom_fn, - invalid_multimodal_strategy=invalid_multimodal_strategy, - apply_chat_template_kwargs=getattr(self.config, "apply_chat_template_kwargs", None), - ) + if is_preference_mode(self.config): + self._eval_dataset = PreferenceStreamingDataset( + path=[d.path for d in eval_prompt_data], + tokenizer=self._tokenizer, + prompt_key=eval_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, + source_name="+".join(d.name for d in eval_prompt_data), + max_length=self.config.preference_max_length, + max_completion_length=self.config.preference_max_completion_length, + pair_capacity=capacity, + seed=seed, + prefetch_max_cached=0, + 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, + ) + else: + self._eval_dataset = SFTStreamingDataset( + path=[d.path for d in eval_prompt_data], + tokenizer=self._tokenizer, + processor_pool=self._processor_pool, + capacity=capacity, + prompt_key=eval_input_key, + label_key=eval_label_key, + multimodal_keys=self.config.multimodal_keys, + conversation_key_map=getattr(self.config, "conversation_key_map", None), + metadata_key=self.config.metadata_key, + tool_key=eval_tool_key, + system_prompt=self.config.system_prompt, + source_name="+".join(d.name for d in eval_prompt_data), + seed=seed, + prefetch_max_cached=0, + pad_token_ids=pad_token_ids, + oversize_strategy=oversize_strategy, + oversize_custom_fn=oversize_custom_fn, + invalid_multimodal_strategy=invalid_multimodal_strategy, + apply_chat_template_kwargs=getattr(self.config, "apply_chat_template_kwargs", None), + ) # Resume: align IndexManager with `start_rollout_id` so a restart sees # the same shuffled order it would on a fresh run. @@ -300,10 +347,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,17 +386,28 @@ 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})" ) - await self._maybe_produce_eval() + await self._maybe_produce_eval(self.step + 1) self.step += 1 def _build_eval_batches(self) -> list[ProcessedSample] | None: @@ -358,8 +422,8 @@ def _build_eval_batches(self) -> list[ProcessedSample] | None: return self._eval_dataset.get_batch_in_order(0, len(self._eval_dataset)) return None - async def _maybe_produce_eval(self) -> None: - """Push the eval set under partitions ``sft_eval__n_`` when + async def _maybe_produce_eval(self, completed_steps: int) -> None: + """Push eval partitions keyed by completed optimizer-step count when due, chunked into ``global_batch_size`` pieces and serially drained. TQ per-partition storage is sized for ``global_batch_size`` (one train @@ -369,66 +433,103 @@ async def _maybe_produce_eval(self) -> None: discover it, and push-then-wait-for-drain serially. Eval blocks the producer here, but only on eval steps. """ - eval_interval = getattr(self.config, "eval_interval", None) - if not eval_interval or eval_interval <= 0: - return - if (self.step + 1) % eval_interval != 0: + if not should_run_sft_eval(self.config, completed_steps): return samples = self._build_eval_batches() if samples is None: return if not samples: raise RuntimeError( - f"Eval @ step {self.step}: source produced 0 valid samples. Refusing to skip the eval push because " + f"Eval @ completed_steps={completed_steps}: source produced 0 valid samples. " + "Refusing to skip the eval push because " "the Megatron consumer is waiting for an eval partition. Check invalid-multimodal and oversize " "skip warnings." ) + if is_preference_mode(self.config): + from relax.engine.sft.eval.acceptance import ( + PREFERENCE_PROBE_PAIR_COUNT, + preference_eval_chunk_sizes, + preference_eval_local_batch_sizes, + record_probe_contract, + ) + + samples = sorted(samples, key=lambda pair: pair.pair_id) + record_probe_contract( + getattr(self.config, "save", None), + self.config.sft_objective, + completed_steps, + samples, + ) - # Pad sub-gbs eval pools with random resamples so the eval set always - # forms at least one full ``global_batch_size`` chunk. Without this the - # chunking loop below would skip eval entirely (n_chunks==0), and the - # consumer — which enters ``run_sft_eval`` purely on interval — would - # block forever waiting for partitions that never come. Seeded by step - # so the padding is reproducible across restarts. + # Causal eval pads sub-GBS pools because its legacy consumer requests a + # fixed batch size. Preference eval uses actual partial-chunk sizes and + # must preserve every unique probe pair without padding. gbs = self.config.global_batch_size n_original = len(samples) - if n_original < gbs: - rng = random.Random(self.step) + preference_mode = is_preference_mode(self.config) + if n_original < gbs and not preference_mode: + rng = random.Random(completed_steps) pad_count = gbs - n_original samples = list(samples) + rng.choices(samples, k=pad_count) self._logger.warning( - f"Eval @ step {self.step}: eval pool of {n_original} samples is smaller than " + f"Eval @ completed_steps={completed_steps}: eval pool of {n_original} samples is smaller than " f"global_batch_size ({gbs}); random-padded with {pad_count} resampled (with " f"replacement) samples to fill one batch. PPL counts duplicated samples — " f"interpret with caution." ) - backend_batch = pack_samples_for_tq(samples, force_multimodal_field=self.config.multimodal_keys is not None) + if preference_mode: + backend_batch, preference_custom_meta = pack_preference_pairs_for_tq(samples) + else: + backend_batch = pack_samples_for_tq( + samples, force_multimodal_field=self.config.multimodal_keys is not None + ) + preference_custom_meta = None assert backend_batch is not None - n_samples = len(backend_batch["tokens"]) + row_key = "pair_ids" if is_preference_mode(self.config) else "tokens" + n_samples = len(backend_batch[row_key]) + if preference_mode and n_samples != PREFERENCE_PROBE_PAIR_COUNT: + raise RuntimeError( + "preference eval packing must preserve exactly " + f"{PREFERENCE_PROBE_PAIR_COUNT} probe pairs, got {n_samples}" + ) # Drain the current train partition so the eval chunks have the full # TQ capacity to themselves. - await self._wait_for_partition_drained(f"sft_{self.step}") + if completed_steps > 0: + await self._wait_for_partition_drained(f"sft_{completed_steps - 1}") chunk_size = self.config.global_batch_size - # Drop trailing samples that don't fill a full chunk. The consumer's + # Causal eval retains the legacy full-chunk requirement. Preference + # eval uses actual per-chunk batch sizes below and never drops pairs. + # The causal consumer's # `_get_data_from_transfer_queue` calls `tq.get_meta(batch_size=...)` # which returns size=0 when the partition has fewer than batch_size # samples, so a partial last chunk would never be marked consumed and # the actor's `while not all_consumed` loop would spin forever (it # already burned a full eval round in the wild — see the # `[get_data_profile] samples=0` log spam). - n_chunks = n_samples // chunk_size - n_dropped = n_samples - n_chunks * chunk_size + if preference_mode: + chunk_sizes = preference_eval_chunk_sizes(n_samples, chunk_size) + preference_eval_local_batch_sizes( + n_samples, + chunk_size, + int(getattr(self.config, "data_parallel_size", 1)), + ) + n_chunks = len(chunk_sizes) + n_dropped = 0 + else: + n_chunks = n_samples // chunk_size + n_dropped = n_samples - n_chunks * chunk_size + chunk_sizes = [chunk_size] * n_chunks if n_chunks == 0: raise RuntimeError( - f"Eval @ step {self.step}: eval pool of {n_samples} samples is smaller than " + f"Eval @ completed_steps={completed_steps}: eval pool of {n_samples} samples is smaller than " f"global_batch_size ({chunk_size}); cannot push the full partition expected by the consumer." ) if n_dropped > 0: self._logger.warning( - f"Eval @ step {self.step}: dropping {n_dropped} trailing sample(s) so eval " + f"Eval @ completed_steps={completed_steps}: dropping {n_dropped} trailing sample(s) so eval " f"chunks align to global_batch_size ({chunk_size}); raise eval pool size or " f"reduce global_batch_size if this matters." ) @@ -438,21 +539,26 @@ async def _maybe_produce_eval(self) -> None: # step. On timeout we clear our own pending chunk and bail. chunk_drain_timeout = float(getattr(self.config, "sft_eval_chunk_drain_timeout_sec", 600.0)) self._logger.info( - f"Eval @ step {self.step}: pushing {n_chunks * chunk_size} samples in {n_chunks} chunk(s) of {chunk_size}." + f"Eval @ completed_steps={completed_steps}: pushing {sum(chunk_sizes)} samples " + f"in {n_chunks} chunk(s) of {chunk_size}." ) - for chunk_idx in range(n_chunks): - s = chunk_idx * chunk_size - e = s + chunk_size + offset = 0 + for chunk_idx, current_chunk_size in enumerate(chunk_sizes): + s = offset + e = s + current_chunk_size + offset = e chunk = {k: v[s:e] for k, v in backend_batch.items()} - partition_id = f"sft_eval_{self.step}_n{n_chunks}_{chunk_idx}" + partition_id = f"sft_eval_{completed_steps}_n{n_chunks}_{chunk_idx}" await self.data_system_client.async_put( - data=dict_to_tensordict(chunk, batch_size=len(chunk["tokens"])), + data=dict_to_tensordict(chunk, batch_size=len(chunk[row_key])), partition_id=partition_id, + custom_meta=None if preference_custom_meta is None else preference_custom_meta[s:e], ) drained = await self._wait_for_partition_drained(partition_id, timeout_sec=chunk_drain_timeout) if not drained: self._logger.warning( - f"Eval @ step {self.step}: chunk {chunk_idx}/{n_chunks} ({partition_id}) did not drain " + f"Eval @ completed_steps={completed_steps}: chunk {chunk_idx}/{n_chunks} " + f"({partition_id}) did not drain " f"within {chunk_drain_timeout}s; aborting eval push and clearing TQ." ) await self.data_system_client.async_clear_partition(partition_id=partition_id) @@ -466,6 +572,8 @@ async def run(self) -> None: async def _async_run(self) -> None: try: + if self.step == 0 and is_preference_mode(self.config): + await self._maybe_produce_eval(0) while self.step < self.config.num_rollout and not self._stop_event.is_set(): await self._produce_one_step() except Exception as exc: diff --git a/relax/engine/sft/bootstrap.py b/relax/engine/sft/bootstrap.py index 5871bde28..bd953f20a 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") in {"dpo", "reward_model"}: + 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..886db4cbc --- /dev/null +++ b/relax/engine/sft/dataset/preference.py @@ -0,0 +1,540 @@ +# 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 + chosen_score_position: int + rejected_score_position: 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(), + chosen_score_position=chosen_tokens.numel() - 1, + rejected_score_position=rejected_tokens.numel() - 1, + 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") + from relax.engine.sft.eval.acceptance import encoded_pair_id + + encoded_pair_ids = [encoded_pair_id(pair.pair_id) 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], + "chosen_score_positions": [pair.chosen_score_position for pair in pairs], + "rejected_score_positions": [pair.rejected_score_position 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/eval/acceptance.py b/relax/engine/sft/eval/acceptance.py new file mode 100644 index 000000000..61a9a1a90 --- /dev/null +++ b/relax/engine/sft/eval/acceptance.py @@ -0,0 +1,205 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Reproducible preference-evaluation artifacts and RFC statistics.""" + +import hashlib +import json +from pathlib import Path +from typing import Any, Sequence + + +PREFERENCE_PROBE_PAIR_COUNT = 512 + + +def preference_eval_chunk_sizes(pair_count: int, global_batch_size: int) -> list[int]: + """Split every real probe pair into capacity-bounded eval chunks.""" + if pair_count <= 0 or global_batch_size <= 0: + raise ValueError("preference eval pair count and global batch size must be positive") + full_chunks, remainder = divmod(pair_count, global_batch_size) + return [global_batch_size] * full_chunks + ([remainder] if remainder else []) + + +def preference_eval_local_batch_sizes(pair_count: int, global_batch_size: int, dp_size: int) -> list[int]: + if dp_size <= 0: + raise ValueError("preference eval data-parallel size must be positive") + chunk_sizes = preference_eval_chunk_sizes(pair_count, global_batch_size) + invalid = [size for size in chunk_sizes if size % dp_size != 0] + if invalid: + raise ValueError( + f"preference eval chunk sizes must be divisible by data-parallel size: chunks={invalid}, dp={dp_size}" + ) + return [size // dp_size for size in chunk_sizes] + + +def encoded_pair_id(pair_id: str) -> int: + return int.from_bytes(hashlib.sha256(pair_id.encode()).digest()[:8], "big") >> 1 + + +def canonical_sha256(value: Any) -> str: + payload = json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(payload).hexdigest() + + +def paired_bootstrap(values: Sequence[float], *, seed: int = 1234, num_replicates: int = 10_000) -> dict[str, Any]: + """Compute the RFC's one-sided percentile bootstrap without dtype drift.""" + import numpy as np + + array = np.asarray(values, dtype=np.float64) + if array.ndim != 1 or array.size == 0: + raise ValueError("paired bootstrap requires a non-empty one-dimensional vector") + rng = np.random.Generator(np.random.PCG64(seed)) + indices = rng.integers(0, array.size, size=(num_replicates, array.size), endpoint=False) + replicates = array[indices].mean(axis=1) + if replicates.dtype != np.float64: + raise RuntimeError(f"bootstrap replicates must remain float64, got {replicates.dtype}") + lower_95 = np.quantile(replicates, 0.05, method="linear") + return { + "numpy_version": np.__version__, + "seed": seed, + "num_replicates": num_replicates, + "point_estimate": float(array.mean()), + "lower_95": float(lower_95), + "passes_lower_bound_gt_0_50": bool(lower_95 > 0.50), + "indices_sha256": hashlib.sha256(indices.astype(" Path | None: + return None if not save_path else Path(save_path) / "preference_eval" + + +def _atomic_write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_suffix(path.suffix + ".tmp") + temporary.write_text(json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8") + temporary.replace(path) + + +def record_probe_contract( + save_path: str | None, + objective: str, + completed_steps: int, + pairs, + *, + expected_pair_count: int = PREFERENCE_PROBE_PAIR_COUNT, +) -> dict[str, Any] | None: + if len(pairs) != expected_pair_count: + raise RuntimeError(f"preference eval requires exactly {expected_pair_count} probe pairs, got {len(pairs)}") + directory = artifact_directory(save_path) + if directory is None: + return None + rows = [ + { + "pair_id": pair.pair_id, + "encoded_pair_id": encoded_pair_id(pair.pair_id), + "chosen_tokens": pair.chosen_tokens.tolist(), + "rejected_tokens": pair.rejected_tokens.tolist(), + "chosen_raw_loss_mask": pair.chosen_loss_mask.tolist(), + "rejected_raw_loss_mask": pair.rejected_loss_mask.tolist(), + "chosen_score_position": pair.chosen_score_position, + "rejected_score_position": pair.rejected_score_position, + } + for pair in pairs + ] + contract = { + "objective": objective, + "pair_count": len(rows), + "pair_ids": [row["pair_id"] for row in rows], + "encoded_pair_id_map": {str(row["encoded_pair_id"]): row["pair_id"] for row in rows}, + "probe_sha256": canonical_sha256(rows), + } + baseline_path = directory / f"{objective}-probe-contract.json" + if completed_steps == 0: + _atomic_write_json(baseline_path, contract) + else: + if not baseline_path.is_file(): + raise RuntimeError(f"preference final eval is missing step-0 probe contract: {baseline_path}") + baseline = json.loads(baseline_path.read_text(encoding="utf-8")) + if contract != baseline: + raise RuntimeError( + "preference probe preprocessing/order changed between step 0 and final eval: " + f"baseline={baseline.get('probe_sha256')}, current={contract['probe_sha256']}" + ) + return contract + + +def write_pair_artifacts( + save_path: str | None, + objective: str, + completed_steps: int, + encoded_rows: list[dict[str, Any]], + batch_plan: list[dict[str, Any]], +) -> dict[str, Any] | None: + directory = artifact_directory(save_path) + if directory is None: + return None + contract_path = directory / f"{objective}-probe-contract.json" + if not contract_path.is_file(): + raise RuntimeError(f"preference eval is missing probe contract: {contract_path}") + contract = json.loads(contract_path.read_text(encoding="utf-8")) + pair_id_map = contract["encoded_pair_id_map"] + rows = [] + for row in encoded_rows: + row = dict(row) + encoded = str(row.pop("encoded_pair_id")) + if encoded not in pair_id_map: + raise RuntimeError(f"preference eval produced unknown encoded pair id {encoded}") + row["pair_id"] = pair_id_map[encoded] + rows.append(row) + rows.sort(key=lambda row: contract["pair_ids"].index(row["pair_id"])) + if len(rows) != contract["pair_count"] or len({row["pair_id"] for row in rows}) != len(rows): + raise RuntimeError( + "preference eval pair artifact is incomplete or duplicated: " + f"expected={contract['pair_count']}, actual={len(rows)}" + ) + + plan_sha256 = canonical_sha256(batch_plan) + baseline_summary_path = directory / f"{objective}-step-0000000-summary.json" + if completed_steps != 0: + if not baseline_summary_path.is_file(): + raise RuntimeError("final preference eval is missing step-0 batch-plan evidence") + baseline = json.loads(baseline_summary_path.read_text(encoding="utf-8")) + if plan_sha256 != baseline["batch_plan_sha256"]: + raise RuntimeError( + "preference eval batch plan changed between step 0 and final: " + f"baseline={baseline['batch_plan_sha256']}, current={plan_sha256}" + ) + + if objective == "reward_model": + accuracy_values = [float(row["chosen_score"] > row["rejected_score"]) for row in rows] + else: + epsilon = 1e-6 + accuracy_values = [ + 1.0 if row["reward_margin"] > epsilon else 0.0 if row["reward_margin"] < -epsilon else 0.5 for row in rows + ] + summary = { + "objective": objective, + "completed_steps": completed_steps, + "pair_count": len(rows), + "probe_sha256": contract["probe_sha256"], + "batch_plan_sha256": plan_sha256, + "bootstrap": paired_bootstrap(accuracy_values), + } + stem = f"{objective}-step-{completed_steps:07d}" + rows_path = directory / f"{stem}-pairs.jsonl" + rows_path.parent.mkdir(parents=True, exist_ok=True) + rows_path.write_text( + "".join(json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n" for row in rows), + encoding="utf-8", + ) + _atomic_write_json(directory / f"{stem}-batch-plan.json", batch_plan) + _atomic_write_json(directory / f"{stem}-summary.json", summary) + return summary + + +__all__ = [ + "PREFERENCE_PROBE_PAIR_COUNT", + "canonical_sha256", + "encoded_pair_id", + "paired_bootstrap", + "preference_eval_chunk_sizes", + "preference_eval_local_batch_sizes", + "record_probe_contract", + "write_pair_artifacts", +] diff --git a/relax/engine/sft/eval/preference.py b/relax/engine/sft/eval/preference.py new file mode 100644 index 000000000..74eb7ac93 --- /dev/null +++ b/relax/engine/sft/eval/preference.py @@ -0,0 +1,94 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Forward-only callbacks and reducers for held-out preference evaluation.""" + +import torch + +from relax.utils.training.preference_utils import select_packed_sequence_scores + + +def extract_preference_eval_pair_ids(pair_data: dict) -> list[int]: + """Read one ID per pair from either expanded or raw preference data.""" + pair_ids = pair_data.get("preference_pair_ids") + if pair_ids is None: + pair_ids = pair_data.get("pair_ids") + if pair_ids is None: + raise RuntimeError("preference eval data is missing pair IDs; expected preference_pair_ids or pair_ids") + return [int(value) for value in pair_ids] + + +def compute_reward_model_eval_step( + logits: torch.Tensor, + *, + total_lengths, + score_positions=None, + raw_loss_masks=None, + packed_tokens=None, + branch_tokens=None, + cu_seqlens=None, + **_, +) -> tuple[torch.Tensor, dict[str, list[torch.Tensor]]]: + if score_positions is None: + raise ValueError("reward-model eval batch is missing score_positions") + scores = select_packed_sequence_scores( + logits, + total_lengths, + score_positions, + raw_loss_masks=raw_loss_masks, + packed_tokens=packed_tokens, + branch_tokens=branch_tokens, + cu_seqlens=cu_seqlens, + ).detach() + # Emit one scalar per branch so ``forward_only`` can restore original + # sample order after dynamic micro-batch length balancing. + return torch.empty((0,), device=logits.device), {"scores": list(scores.unbind())} + + +def pair_metric_sums( + chosen: torch.Tensor, + rejected: torch.Tensor, + losses: torch.Tensor, + *, + epsilon: float = 1e-6, +) -> torch.Tensor: + """Return additive loss/count/score/margin/correct/tie statistics.""" + if chosen.shape != rejected.shape or chosen.shape != losses.shape: + raise ValueError("preference eval tensors must have identical shapes") + margin = chosen - rejected + correct = (margin > epsilon).to(torch.float64).sum() + ties = (margin.abs() <= epsilon).to(torch.float64).sum() + return torch.stack( + [ + losses.to(torch.float64).sum(), + torch.tensor(float(losses.numel()), device=losses.device, dtype=torch.float64), + chosen.to(torch.float64).sum(), + rejected.to(torch.float64).sum(), + margin.to(torch.float64).sum(), + correct, + ties, + ] + ) + + +def finalize_pair_metrics(values: torch.Tensor, *, prefix: str) -> dict[str, float]: + loss_sum, count, chosen_sum, rejected_sum, margin_sum, correct, ties = values.tolist() + if count <= 0: + raise ValueError("preference evaluator received zero pairs") + return { + f"eval/{prefix}_loss": loss_sum / count, + f"eval/{prefix}_chosen": chosen_sum / count, + f"eval/{prefix}_rejected": rejected_sum / count, + f"eval/{prefix}_margin": margin_sum / count, + f"eval/{prefix}_strict_accuracy": correct / count, + f"eval/{prefix}_tie_rate": ties / count, + f"eval/{prefix}_tie_aware_accuracy": (correct + 0.5 * ties) / count, + f"eval/{prefix}_pairs": count, + } + + +__all__ = [ + "compute_reward_model_eval_step", + "extract_preference_eval_pair_ids", + "finalize_pair_metrics", + "pair_metric_sums", +] diff --git a/relax/engine/sft/eval/runner.py b/relax/engine/sft/eval/runner.py index 189cf16ec..c2e6217e1 100644 --- a/relax/engine/sft/eval/runner.py +++ b/relax/engine/sft/eval/runner.py @@ -98,6 +98,12 @@ def run_sft_eval(actor, rollout_id: int) -> None: loss mask), and the final all-reduce below includes the CP group so each rank ends up with the full-sequence totals. """ + from relax.engine.sft.runtime import is_preference_mode + + if is_preference_mode(actor.args): + _run_preference_eval(actor, rollout_id) + return + # Lazy imports: keep this module importable without Megatron initialized. from relax.backends.megatron.data import get_data_iterator from relax.backends.megatron.initialize import is_megatron_main_rank @@ -182,3 +188,209 @@ def run_sft_eval(actor, rollout_id: int) -> None: # step) and never reach ClearML/W&B/TB. tracking_utils.flush_metrics(args, step) logger.info(f"SFT eval @ rollout_id={rollout_id}: {metrics}") + + +def _run_preference_eval(actor, rollout_id: int) -> None: + """Evaluate DPO or RM on pair rows using the same TQ packing as + training.""" + from relax.backends.megatron.data import expand_preference_rollout_data, get_data_iterator + from relax.backends.megatron.initialize import is_megatron_main_rank + from relax.backends.megatron.model import forward_only + from relax.engine.sft.eval.acceptance import ( + PREFERENCE_PROBE_PAIR_COUNT, + preference_eval_chunk_sizes, + preference_eval_local_batch_sizes, + ) + from relax.engine.sft.eval.preference import ( + compute_reward_model_eval_step, + extract_preference_eval_pair_ids, + finalize_pair_metrics, + pair_metric_sums, + ) + from relax.utils.training.preference_utils import dpo_pair_loss, reward_model_pair_loss + + args = actor.args + task_name = "sft_eval" + dp_size = mpu.get_data_parallel_world_size(with_context_parallel=False) + data_fields = [ + "pair_ids", + "chosen_tokens", + "rejected_tokens", + "chosen_loss_masks", + "rejected_loss_masks", + "chosen_total_lengths", + "rejected_total_lengths", + "chosen_score_positions", + "rejected_score_positions", + ] + n_chunks = _wait_for_eval_chunk_count(actor, rollout_id) + chunk_sizes = preference_eval_chunk_sizes(PREFERENCE_PROBE_PAIR_COUNT, args.global_batch_size) + local_batch_sizes = preference_eval_local_batch_sizes(PREFERENCE_PROBE_PAIR_COUNT, args.global_batch_size, dp_size) + if len(chunk_sizes) != n_chunks: + raise RuntimeError(f"preference eval chunk plan mismatch: producer={n_chunks}, consumer={len(chunk_sizes)}") + local = torch.zeros(7, device=device_utils.make_current_torch_device(), dtype=torch.float64) + local_rows: list[dict] = [] + local_plan: list[dict] = [] + started = time.monotonic() + with timer("preference_eval"): + for chunk_idx, (global_chunk_size, batch_size) in enumerate(zip(chunk_sizes, local_batch_sizes, strict=True)): + partition_id = f"sft_eval_{rollout_id}_n{n_chunks}_{chunk_idx}" + _wait_for_eval_partition_present(actor, partition_id) + batch_index = 0 + while not actor.all_consumed(task_name, rollout_id, partition_id=partition_id): + pair_rows, _batch_meta = actor._get_data_from_transfer_queue( + task_name, rollout_id, data_fields, batch_size, batch_index, partition_id=partition_id + ) + if pair_rows is None: + continue + batch_index += 1 + rollout_data = expand_preference_rollout_data(pair_rows) + rollout_data["dynamic_global_batch_size"] = global_chunk_size + data_iterator, num_microbatches = get_data_iterator(args, actor.model, rollout_data) + for microbatch_index, branch_indices in enumerate(data_iterator[0].micro_batch_indices): + encoded_ids = [] + for branch_index in branch_indices: + pair_id = int(rollout_data["preference_branch_pair_ids"][branch_index]) + if not encoded_ids or encoded_ids[-1] != pair_id: + encoded_ids.append(pair_id) + local_plan.append( + { + "rank": dist.get_rank(), + "chunk": chunk_idx, + "batch": batch_index - 1, + "microbatch": microbatch_index, + "encoded_pair_ids": encoded_ids, + } + ) + encoded_pair_ids = extract_preference_eval_pair_ids(rollout_data) + if args.sft_objective == "dpo": + if args.dpo_reference_free: + reference_sums = None + else: + actor._switch_model("ref") + reference = actor.compute_log_prob(data_iterator, num_microbatches, store_prefix="ref_")[ + "ref_log_probs" + ] + reference_sums = _masked_sequence_sums(reference, rollout_data["loss_masks"], local.device) + actor._switch_model("actor") + policy = actor.compute_log_prob(data_iterator, num_microbatches, store_prefix="")["log_probs"] + policy_sums = _masked_sequence_sums(policy, rollout_data["loss_masks"], local.device) + policy_chosen, policy_rejected = policy_sums[0::2], policy_sums[1::2] + if reference_sums is None: + reference_chosen = reference_rejected = None + chosen_values = args.dpo_beta * policy_chosen + rejected_values = args.dpo_beta * policy_rejected + else: + reference_chosen, reference_rejected = reference_sums[0::2], reference_sums[1::2] + chosen_values = args.dpo_beta * (policy_chosen - reference_chosen) + rejected_values = args.dpo_beta * (policy_rejected - reference_rejected) + losses = dpo_pair_loss( + policy_chosen, + policy_rejected, + reference_chosen=reference_chosen, + reference_rejected=reference_rejected, + beta=args.dpo_beta, + reference_free=args.dpo_reference_free, + ) + local += pair_metric_sums(chosen_values, rejected_values, losses) + policy_chosen_values = policy_chosen.detach().cpu().tolist() + policy_rejected_values = policy_rejected.detach().cpu().tolist() + reference_chosen_values = ( + [0.0] * len(encoded_pair_ids) + if reference_chosen is None + else reference_chosen.detach().cpu().tolist() + ) + reference_rejected_values = ( + [0.0] * len(encoded_pair_ids) + if reference_rejected is None + else reference_rejected.detach().cpu().tolist() + ) + chosen_reward_values = chosen_values.detach().cpu().tolist() + rejected_reward_values = rejected_values.detach().cpu().tolist() + loss_values = losses.detach().cpu().tolist() + for index, encoded_pair_id in enumerate(encoded_pair_ids): + margin = chosen_reward_values[index] - rejected_reward_values[index] + local_rows.append( + { + "encoded_pair_id": encoded_pair_id, + "policy_chosen_logp": policy_chosen_values[index], + "policy_rejected_logp": policy_rejected_values[index], + "reference_chosen_logp": reference_chosen_values[index], + "reference_rejected_logp": reference_rejected_values[index], + "chosen_implicit_reward": chosen_reward_values[index], + "rejected_implicit_reward": rejected_reward_values[index], + "reward_margin": margin, + "pair_loss": loss_values[index], + } + ) + else: + outputs = forward_only( + compute_reward_model_eval_step, + args, + actor.model, + data_iterator, + num_microbatches, + store_prefix="", + ) + if mpu.is_pipeline_last_stage(): + scores = torch.stack(outputs["scores"]).to(local.device) + chosen_scores, rejected_scores = scores[0::2], scores[1::2] + losses = reward_model_pair_loss(chosen_scores, rejected_scores) + local += pair_metric_sums(chosen_scores, rejected_scores, losses, epsilon=0.0) + chosen_values = chosen_scores.detach().cpu().tolist() + rejected_values = rejected_scores.detach().cpu().tolist() + loss_values = losses.detach().cpu().tolist() + for index, encoded_pair_id in enumerate(encoded_pair_ids): + local_rows.append( + { + "encoded_pair_id": encoded_pair_id, + "chosen_score": chosen_values[index], + "rejected_score": rejected_values[index], + "pair_loss": loss_values[index], + } + ) + dist.barrier(group=get_gloo_group()) + if dist.get_rank() == 0: + run(actor.data_system_client.async_clear_partition(partition_id=partition_id)) + + dist.all_reduce(local, op=dist.ReduceOp.SUM, group=mpu.get_pipeline_model_parallel_group()) + dist.all_reduce(local, op=dist.ReduceOp.SUM, group=mpu.get_data_parallel_group(with_context_parallel=True)) + metrics = finalize_pair_metrics(local, prefix="dpo" if args.sft_objective == "dpo" else "rm") + metrics["perf/preference_eval_time"] = time.monotonic() - started + gloo_group = get_gloo_group() + gathered_rows = [None] * dist.get_world_size(group=gloo_group) + gathered_plans = [None] * dist.get_world_size(group=gloo_group) + dist.all_gather_object(gathered_rows, local_rows, group=gloo_group) + dist.all_gather_object(gathered_plans, local_plan, group=gloo_group) + if is_megatron_main_rank(): + from relax.engine.sft.eval.acceptance import write_pair_artifacts + + summary = write_pair_artifacts( + getattr(args, "save", None), + args.sft_objective, + rollout_id, + [row for rank_rows in gathered_rows for row in rank_rows], + [entry for rank_plan in gathered_plans for entry in rank_plan], + ) + if summary is not None: + metrics[f"eval/{'dpo' if args.sft_objective == 'dpo' else 'rm'}_bootstrap_lower_95"] = summary[ + "bootstrap" + ]["lower_95"] + step = compute_rollout_step(args, rollout_id) + metrics["rollout/step"] = step + tracking_utils.log(args, metrics, step_key="rollout/step") + tracking_utils.flush_metrics(args, step) + logger.info(f"Preference eval @ rollout_id={rollout_id}: {metrics}") + + +def _masked_sequence_sums(values, masks, device: torch.device) -> torch.Tensor: + if len(values) != len(masks): + raise ValueError("preference eval values/masks are not branch aligned") + sums = [] + for value, mask in zip(values, masks, strict=True): + value = torch.as_tensor(value, device=device) + mask = torch.as_tensor(mask, device=device, dtype=value.dtype) + if value.shape != mask.shape: + raise ValueError(f"preference eval value/mask shape mismatch: {value.shape} vs {mask.shape}") + sums.append((value * mask).sum()) + return torch.stack(sums) diff --git a/relax/engine/sft/runtime.py b/relax/engine/sft/runtime.py index d5c1a51cc..6b2e8b895 100644 --- a/relax/engine/sft/runtime.py +++ b/relax/engine/sft/runtime.py @@ -7,6 +7,8 @@ here keeps the dispatchers in those files to one-line calls. """ +import math +import re from argparse import Namespace @@ -19,6 +21,96 @@ 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) in {"dpo", "reward_model"} + + +def validate_preference_args(args: Namespace) -> None: + """Reject unsupported preference configurations before Serve starts.""" + if not is_preference_mode(args): + return + objective = sft_objective(args) + if objective == "reward_model" and getattr(args, "save_hf", None) is not None: + raise ValueError( + "reward_model v1 does not support --save-hf; use native Megatron checkpoints for RM persistence" + ) + 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 objective != "dpo" and getattr(args, "dpo_reference_free", False): + raise ValueError("--dpo-reference-free is valid only with --sft-objective dpo") + if objective == "dpo": + 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( + "DPO requires --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( + "DPO 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}" @@ -37,8 +129,14 @@ def sft_task_name(args: Namespace, *, component: str = "actor") -> str: return "train" -def should_run_sft_eval(args: Namespace, rollout_id: int) -> bool: - """SFT PPL eval triggers every ``--eval-interval`` steps under SFT mode +def should_run_sft_eval(args: Namespace, completed_steps: int) -> bool: + """Return whether eval is due after ``completed_steps`` optimizer steps. + + Preference objectives additionally evaluate at the true pre-training + baseline (0) and at the final completed step, independent of whether the + periodic interval happens to divide the run length. + + SFT PPL eval triggers every ``--eval-interval`` steps under SFT mode when an eval source is configured (either ``--eval-prompt-data`` or ``--eval-size``, mutually exclusive — see ``utils/arguments.py``). @@ -52,11 +150,28 @@ def should_run_sft_eval(args: Namespace, rollout_id: int) -> bool: interval = getattr(args, "eval_interval", None) if interval is None or interval <= 0: return False - return (rollout_id + 1) % interval == 0 + if is_preference_mode(args) and completed_steps in {0, int(getattr(args, "num_rollout", 0) or 0)}: + return True + return completed_steps > 0 and completed_steps % interval == 0 -def should_run_sft_predict(args: Namespace, rollout_id: int) -> bool: - """SFT periodic predict triggers every ``--sft-predict-interval`` steps. +def actor_training_input_ready(args: Namespace, step: int, partition_ids: list[str] | None) -> bool: + """Allow the backend to enter step 0 when its preference baseline is ready. + + In colocate mode the Actor service normally waits for the train partition + before calling the backend. Preference producers intentionally publish and + drain the step-0 eval partition first, so that baseline partition is the + backend's first input and must also release the service-level wait. + """ + if not partition_ids: + return False + if step == 0 and is_preference_mode(args) and should_run_sft_eval(args, completed_steps=0): + return any(re.fullmatch(r"sft_eval_0_n\d+_0", partition_id) for partition_id in partition_ids) + return sft_partition_id(args, step) in partition_ids + + +def should_run_sft_predict(args: Namespace, completed_steps: int) -> bool: + """SFT periodic predict triggers after each completed interval. Argparse already validated ``--loss-type sft``, ``--save``, and the eval data source, so we only need the interval check here. @@ -64,4 +179,9 @@ def should_run_sft_predict(args: Namespace, rollout_id: int) -> bool: interval = getattr(args, "sft_predict_interval", None) if interval is None or interval <= 0: return False - return (rollout_id + 1) % interval == 0 + return completed_steps > 0 and completed_steps % interval == 0 + + +def evaluation_step_for_rollout(args: Namespace, rollout_id: int) -> int: + """Map a zero-based training rollout to the evaluation step namespace.""" + return rollout_id + 1 if is_sft_mode(args) else rollout_id diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index f7d1e710e..34cc604ca 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", "reward_model"], + 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..752a0cc62 100644 --- a/relax/utils/training/data_fields.py +++ b/relax/utils/training/data_fields.py @@ -27,6 +27,18 @@ 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") in {"dpo", "reward_model"}: + return [ + "pair_ids", + "chosen_tokens", + "rejected_tokens", + "chosen_loss_masks", + "rejected_loss_masks", + "chosen_total_lengths", + "rejected_total_lengths", + "chosen_score_positions", + "rejected_score_positions", + ] 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/ppo_utils.py b/relax/utils/training/ppo_utils.py index e4eab1310..5008b955e 100644 --- a/relax/utils/training/ppo_utils.py +++ b/relax/utils/training/ppo_utils.py @@ -1140,9 +1140,45 @@ def install_critic_value_head_in_provider( ) +def install_reward_model_head_in_provider( + model: torch.nn.Module, + args, + role: str, + post_process: bool, + *, + stash_lm_head: bool = False, +) -> None: + """Install the offline reward-model scalar head before DDP/optimizer + construction.""" + if ( + role != "actor" + or not post_process + or getattr(args, "loss_type", None) != "sft" + or getattr(args, "sft_objective", "causal_lm") != "reward_model" + ): + return + + owner = _find_output_layer_owner(model) + if owner is None: + return + + output_layer = owner.output_layer + if isinstance(output_layer, LinearForLastLayer) and output_layer.out_features == 1: + return + + if stash_lm_head: + object.__setattr__(owner, _RELAX_HF_OUTPUT_LAYER_ATTR, output_layer) + owner.output_layer = LinearForLastLayer( + input_size=owner.config.hidden_size, + output_size=1, + config=owner.config, + bias=False, + ) + + @contextlib.contextmanager def use_critic_lm_head_for_hf_load(model): - """Temporarily restore the stashed LM head for HF Bridge weight loading. + """Temporarily restore a stashed LM head for HF Bridge weight loading. Bridge can only convert HF weights against a vocab-sized ``output_layer``; the scalar value head is put back in ``finally`` (asserting the exact same @@ -1169,9 +1205,9 @@ def use_critic_lm_head_for_hf_load(model): for owner, value_head, value_param_ids in reversed(restored_heads): owner.output_layer = value_head object.__delattr__(owner, _RELAX_HF_OUTPUT_LAYER_ATTR) - assert owner.output_layer is value_head, "critic value head object changed during HF checkpoint loading" + assert owner.output_layer is value_head, "scalar head object changed during HF checkpoint loading" assert tuple(id(param) for param in value_head.parameters()) == value_param_ids, ( - "critic value head parameters changed during HF checkpoint loading" + "scalar head parameters changed during HF checkpoint loading" ) @@ -1235,6 +1271,30 @@ def validate_critic_value_head_registration(model, optimizer) -> tuple[int, ...] return tuple(value_head_param_ids) +def validate_reward_model_head_registration(model, optimizer) -> tuple[int, ...]: + """Validate RM scalar-head shape, bias contract, registration, and DDP + ownership.""" + del optimizer + parameter_ids = [] + for model_chunk in model: + owner = _find_output_layer_owner(model_chunk) + if owner is None: + continue + head = owner.output_layer + assert isinstance(head, LinearForLastLayer), ( + f"reward-model output layer must be LinearForLastLayer, got {type(head).__name__}" + ) + assert tuple(head.weight.shape) == (1, owner.config.hidden_size), ( + f"reward-model head weight must have shape (1, {owner.config.hidden_size}), got {tuple(head.weight.shape)}" + ) + assert head.bias is None, "reward-model scalar head must use bias=False" + registered_parameter_ids = {id(parameter) for parameter in model_chunk.parameters()} + assert id(head.weight) in registered_parameter_ids, "reward-model output_layer.weight is not registered" + assert _ddp_owns_param(model_chunk, head.weight), "DDP does not own reward-model output_layer.weight" + parameter_ids.append(id(head.weight)) + return tuple(parameter_ids) + + def snapshot_critic_value_head_state(model) -> dict: """One-scalar-per-param snapshot of value head weights, keyed by chunk+name. diff --git a/relax/utils/training/preference_utils.py b/relax/utils/training/preference_utils.py new file mode 100644 index 000000000..47a86a5ce --- /dev/null +++ b/relax/utils/training/preference_utils.py @@ -0,0 +1,279 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Pure helpers shared by DPO and pairwise reward-model 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 reward_model_pair_loss(chosen_scores: torch.Tensor, rejected_scores: torch.Tensor) -> torch.Tensor: + """Return unreduced Bradley-Terry loss, one value per preference pair.""" + _validate_same_shape("reward-model scores", chosen_scores, rejected_scores) + margins = chosen_scores - rejected_scores + require_tensor_condition( + torch.isfinite(margins).all(), + "reward-model margins must contain only finite values", + ) + return -F.logsigmoid(margins) + + +def select_packed_sequence_scores( + logits: torch.Tensor, + total_lengths: Sequence[int], + score_positions: Sequence[int], + *, + raw_loss_masks: Sequence[torch.Tensor] | None = None, + packed_tokens: torch.Tensor | None = None, + branch_tokens: Sequence[torch.Tensor] | None = None, + cu_seqlens: torch.Tensor | None = None, +) -> torch.Tensor: + """Validate and select one terminal score per CP=1 THD branch.""" + if len(total_lengths) != len(score_positions): + raise ValueError( + f"total_lengths/score_positions length mismatch: {len(total_lengths)} vs {len(score_positions)}" + ) + if logits.ndim == 3 and logits.shape[0] == 1 and logits.shape[-1] == 1: + flat_logits = logits[0, :, 0] + elif logits.ndim == 2 and logits.shape[-1] == 1: + flat_logits = logits[:, 0] + elif logits.ndim == 1: + flat_logits = logits + else: + raise ValueError(f"reward-model logits must have shape [1,T,1], [T,1], or [T], got {tuple(logits.shape)}") + + branch_count = len(total_lengths) + optional_fields = { + "raw_loss_masks": raw_loss_masks, + "branch_tokens": branch_tokens, + } + for name, values in optional_fields.items(): + if values is not None and len(values) != branch_count: + raise ValueError(f"{name} must be branch aligned: expected {branch_count}, got {len(values)}") + if (packed_tokens is None) != (branch_tokens is None): + raise ValueError("packed_tokens and branch_tokens must be provided together") + + if packed_tokens is not None: + if packed_tokens.ndim == 2 and packed_tokens.shape[0] == 1: + flat_packed_tokens = packed_tokens[0] + elif packed_tokens.ndim == 1: + flat_packed_tokens = packed_tokens + else: + raise ValueError(f"packed reward tokens must have shape [1,T] or [T], got {tuple(packed_tokens.shape)}") + if flat_packed_tokens.numel() != flat_logits.numel(): + raise ValueError( + f"packed reward token/logit length mismatch: {flat_packed_tokens.numel()} vs {flat_logits.numel()}" + ) + else: + flat_packed_tokens = None + + offsets: list[int] = [] + cursor = 0 + for index, (length, position) in enumerate(zip(total_lengths, score_positions, strict=True)): + length = int(length) + position = int(position) + if length <= 0: + raise ValueError(f"sequence {index} has non-positive total length {length}") + if not 0 <= position < length: + raise ValueError(f"sequence {index} score position {position} is outside [0, {length})") + packed_index = cursor + position + offsets.append(packed_index) + if raw_loss_masks is not None: + raw_mask = torch.as_tensor(raw_loss_masks[index]) + if raw_mask.ndim != 1 or raw_mask.numel() != length: + raise ValueError( + f"sequence {index} raw loss mask must have length {length}, got {tuple(raw_mask.shape)}" + ) + require_tensor_condition( + raw_mask[position] == 1, + f"sequence {index} score position must be supervised by raw_loss_mask", + ) + if flat_packed_tokens is not None and branch_tokens is not None: + branch = torch.as_tensor(branch_tokens[index], device=flat_packed_tokens.device) + if branch.ndim != 1 or branch.numel() != length: + raise ValueError( + f"sequence {index} branch tokens must have length {length}, got {tuple(branch.shape)}" + ) + require_tensor_condition( + flat_packed_tokens[packed_index] == branch[position], + f"sequence {index} packed terminal token does not match branch terminal token", + ) + cursor += length + if cursor > flat_logits.numel(): + raise ValueError(f"packed reward logits contain {flat_logits.numel()} tokens, expected at least {cursor}") + if cu_seqlens is not None: + if cu_seqlens.ndim != 1 or cu_seqlens.numel() not in {branch_count + 1, branch_count + 2}: + raise ValueError( + "reward-model cu_seqlens must describe the real branches and at most one trailing padding segment" + ) + expected = torch.tensor( + [0, *torch.tensor(total_lengths, dtype=torch.long).cumsum(0).tolist()], + device=cu_seqlens.device, + dtype=cu_seqlens.dtype, + ) + require_tensor_condition( + (cu_seqlens[: branch_count + 1] == expected).all(), + "reward-model cu_seqlens do not match branch lengths/order", + ) + if cu_seqlens.numel() == branch_count + 1: + if flat_logits.numel() != cursor: + raise ValueError("reward-model packed tail must be represented by an explicit padding segment") + else: + require_tensor_condition( + cu_seqlens[-1] == flat_logits.numel(), + "reward-model padding segment must cover only the packed tail", + ) + if not offsets: + return flat_logits.new_empty((0,)) + return flat_logits[torch.tensor(offsets, device=flat_logits.device, dtype=torch.long)] + + +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", + "reward_model_pair_loss", + "select_packed_sequence_scores", +] 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..b83a5553c --- /dev/null +++ b/scripts/training/dpo/run-qwen3-0.6B-ultrafeedback-1xgpu.sh @@ -0,0 +1,68 @@ +#!/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}" +EVAL_PROMPT_DATA="${EVAL_PROMPT_DATA:?set EVAL_PROMPT_DATA to the Task 31 held-out 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}" \ + --eval-prompt-data task31 "${EVAL_PROMPT_DATA}" \ + --eval-interval "${EVAL_INTERVAL:-200}" \ + --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/scripts/training/reward_modeling/run-qwen3-0.6B-ultrafeedback-1xgpu.sh b/scripts/training/reward_modeling/run-qwen3-0.6B-ultrafeedback-1xgpu.sh new file mode 100644 index 000000000..feb6af587 --- /dev/null +++ b/scripts/training/reward_modeling/run-qwen3-0.6B-ultrafeedback-1xgpu.sh @@ -0,0 +1,64 @@ +#!/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}" +EVAL_PROMPT_DATA="${EVAL_PROMPT_DATA:?set EVAL_PROMPT_DATA to the Task 31 held-out JSONL or Parquet}" +SAVE_DIR="${SAVE_DIR:-${SCRIPT_DIR}/../../../checkpoints/task31-reward-modeling}" +EXP_NAME="${EXP_NAME:-qwen3-0.6b-ultrafeedback-rm-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 reward_model \ + --prompt-data "${PROMPT_DATA}" \ + --eval-prompt-data task31 "${EVAL_PROMPT_DATA}" \ + --eval-interval "${EVAL_INTERVAL:-200}" \ + --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}" \ + --ref-load "${HF_CHECKPOINT}" \ + --megatron-to-hf-mode bridge \ + --save "${SAVE_DIR}/${EXP_NAME}" \ + --load "${SAVE_DIR}/${EXP_NAME}" \ + --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:-1e-5}" \ + --lr-decay-style cosine \ + --min-lr 0 \ + --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_model_provider_vpp.py b/tests/backends/megatron/test_model_provider_vpp.py index 82c20510b..b4c65613b 100644 --- a/tests/backends/megatron/test_model_provider_vpp.py +++ b/tests/backends/megatron/test_model_provider_vpp.py @@ -211,6 +211,37 @@ def __init__(self): assert all("relax_hf_output_layer" not in name for name, _ in model.named_parameters()) +def test_bridge_reward_model_provider_registers_biasless_scalar_head_and_restores_on_error(monkeypatch): + class _FakeBridgeModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(hidden_size=4, sequence_parallel=False) + self.output_layer = torch.nn.Linear(4, 8) + + module, _ = _load_model_provider(monkeypatch, provider=_FakeProvider(_FakeBridgeModel())) + args = _bridge_args(loss_type="sft", sft_objective="reward_model") + model = module.get_model_provider_func(args, role="actor")(post_process=True) + reward_head = model.output_layer + parameter_ids = tuple(id(parameter) for parameter in reward_head.parameters()) + + assert isinstance(reward_head, ppo_utils.LinearForLastLayer) + assert tuple(reward_head.weight.shape) == (1, 4) + assert reward_head.bias is None + assert list(model.state_dict()) == ["output_layer.weight"] + + with pytest.raises(RuntimeError, match="bridge failed"): + with ppo_utils.use_critic_lm_head_for_hf_load([model]): + assert model.output_layer.out_features == 8 + raise RuntimeError("bridge failed") + + assert model.output_layer is reward_head + assert tuple(id(parameter) for parameter in reward_head.parameters()) == parameter_ids + assert not hasattr(model, ppo_utils._RELAX_HF_OUTPUT_LAYER_ATTR) + + reward_head.weight.main_grad = torch.zeros_like(reward_head.weight) + assert ppo_utils.validate_reward_model_head_registration([model], object()) == parameter_ids + + def test_hf_load_context_restores_same_value_head(monkeypatch): class _FakeBridgeModel(torch.nn.Module): def __init__(self): diff --git a/tests/backends/megatron/test_preference_batching.py b/tests/backends/megatron/test_preference_batching.py new file mode 100644 index 000000000..44f3101dc --- /dev/null +++ b/tests/backends/megatron/test_preference_batching.py @@ -0,0 +1,174 @@ +# 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, + "chosen_score_positions": [1] * count, + "rejected_score_positions": [1] * 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_reward_model_checkpoint.py b/tests/backends/megatron/test_reward_model_checkpoint.py new file mode 100644 index 000000000..5cd367d07 --- /dev/null +++ b/tests/backends/megatron/test_reward_model_checkpoint.py @@ -0,0 +1,323 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Reward-model checkpoint metadata and exact head-schema tests.""" + +import copy +from types import SimpleNamespace + +import pytest +import torch + + +try: + from relax.backends.megatron import checkpoint as checkpoint_module +except Exception as exc: + pytest.skip(f"Megatron checkpoint helpers unavailable: {exc}", allow_module_level=True) + + +class _TensorMetadata: + def __init__(self, shape): + self.global_shape = shape + + +@pytest.mark.parametrize( + ("tracker_value", "expected_directory"), + [("0", "iter_0000000"), ("17", "iter_0000017"), ("release", "release")], +) +def test_checkpoint_iteration_dir_supports_megatron_tracker_formats(tmp_path, tracker_value, expected_directory): + (tmp_path / "latest_checkpointed_iteration.txt").write_text(tracker_value, encoding="utf-8") + assert checkpoint_module._checkpoint_iteration_dir(tmp_path) == tmp_path / expected_directory + + +def test_checkpoint_iteration_dir_rejects_invalid_tracker_metadata(tmp_path): + (tmp_path / "latest_checkpointed_iteration.txt").write_text("invalid", encoding="utf-8") + with pytest.raises(RuntimeError, match="cannot resolve Megatron checkpoint iteration"): + checkpoint_module._checkpoint_iteration_dir(tmp_path) + + +def test_checkpoint_iteration_dir_honors_explicit_checkpoint_step_for_iteration_tracker(tmp_path): + (tmp_path / "latest_checkpointed_iteration.txt").write_text("17", encoding="utf-8") + assert checkpoint_module._checkpoint_iteration_dir(tmp_path, ckpt_step=42) == tmp_path / "iter_0000042" + + +def test_checkpoint_iteration_dir_honors_explicit_zero_checkpoint_step(tmp_path): + (tmp_path / "latest_checkpointed_iteration.txt").write_text("17", encoding="utf-8") + assert checkpoint_module._checkpoint_iteration_dir(tmp_path, ckpt_step=0) == tmp_path / "iter_0000000" + + +def test_checkpoint_iteration_dir_keeps_release_when_checkpoint_step_is_set(tmp_path): + (tmp_path / "latest_checkpointed_iteration.txt").write_text("release", encoding="utf-8") + assert checkpoint_module._checkpoint_iteration_dir(tmp_path, ckpt_step=42) == tmp_path / "release" + + +def test_checkpoint_iteration_dir_rejects_negative_iterations(tmp_path): + (tmp_path / "latest_checkpointed_iteration.txt").write_text("-1", encoding="utf-8") + with pytest.raises(RuntimeError, match="must be non-negative"): + checkpoint_module._checkpoint_iteration_dir(tmp_path) + + +def test_reward_model_tensor_metadata_accepts_exact_bias_free_head(): + checkpoint_module._validate_reward_model_tensor_metadata( + { + "model.output_layer.weight": _TensorMetadata((1, 1024)), + "model.decoder.weight": _TensorMetadata((8, 8)), + "optimizer.state.exp_avg.model.output_layer.weight": _TensorMetadata((1, 1024)), + "optimizer.state.exp_avg_sq.model.output_layer.weight": _TensorMetadata((1, 1024)), + }, + 1024, + ) + + +@pytest.mark.parametrize( + ("metadata", "match"), + [ + ({}, "exactly one"), + ({"model.output_layer.weight": _TensorMetadata((2, 1024))}, "shape mismatch"), + ( + { + "model.output_layer.weight": _TensorMetadata((1, 1024)), + "model.output_layer.bias": _TensorMetadata((1,)), + }, + "unexpected", + ), + ( + { + "model.output_layer.weight": _TensorMetadata((1, 1024)), + "model.reward_model_head.weight": _TensorMetadata((1, 1024)), + }, + "unexpected", + ), + ], +) +def test_reward_model_tensor_metadata_rejects_missing_extra_and_wrong_shape(metadata, match): + with pytest.raises(RuntimeError, match=match): + checkpoint_module._validate_reward_model_tensor_metadata(metadata, 1024) + + +def test_reward_model_contract_rejects_critic_metadata_before_tensor_load(monkeypatch, tmp_path): + calls = [] + fake_dist_checkpointing = SimpleNamespace( + check_is_distributed_checkpoint=lambda path: True, + load_common_state_dict=lambda path: { + "args": SimpleNamespace( + sft_objective="causal_lm", head_type="critic_value_terminal_v1", checkpoint_role="critic" + ) + }, + load_tensors_metadata=lambda path: calls.append(path), + ) + import megatron.core + + monkeypatch.setattr(megatron.core, "dist_checkpointing", fake_dist_checkpointing) + args = SimpleNamespace( + loss_type="sft", + sft_objective="reward_model", + hidden_size=1024, + no_load_optim=False, + no_load_rng=False, + finetune=False, + reset_optimizer_states=False, + ) + model = [SimpleNamespace(role="actor")] + with pytest.raises(RuntimeError, match="RM resume requires checkpoint metadata"): + checkpoint_module._validate_checkpoint_contract(args, model, tmp_path) + assert calls == [] + + +def test_reward_model_contract_rejects_release_checkpoint_before_metadata_load(monkeypatch, tmp_path): + calls = [] + import megatron.core + + monkeypatch.setattr( + megatron.core, + "dist_checkpointing", + SimpleNamespace(load_common_state_dict=lambda path: calls.append(path)), + ) + args = SimpleNamespace(loss_type="sft", sft_objective="reward_model") + with pytest.raises(RuntimeError, match="rejects release checkpoints"): + checkpoint_module._validate_checkpoint_contract(args, [SimpleNamespace(role="actor")], tmp_path / "release") + assert calls == [] + + +@pytest.mark.parametrize( + ("args", "role"), + [ + (SimpleNamespace(loss_type="sft", sft_objective="causal_lm"), "actor"), + (SimpleNamespace(loss_type="sft", sft_objective="dpo"), "actor"), + (SimpleNamespace(), "critic"), + ], +) +def test_non_reward_model_contract_defers_legacy_checkpoint_to_megatron(monkeypatch, tmp_path, args, role): + common_state_loads = [] + fake_dist_checkpointing = SimpleNamespace( + check_is_distributed_checkpoint=lambda path: False, + load_common_state_dict=lambda path: common_state_loads.append(path), + ) + import megatron.core + + monkeypatch.setattr(megatron.core, "dist_checkpointing", fake_dist_checkpointing) + checkpoint_module._validate_checkpoint_contract(args, [SimpleNamespace(role=role)], tmp_path) + assert common_state_loads == [] + + +def test_reward_model_contract_rejects_legacy_checkpoint_before_metadata_load(monkeypatch, tmp_path): + common_state_loads = [] + fake_dist_checkpointing = SimpleNamespace( + check_is_distributed_checkpoint=lambda path: False, + load_common_state_dict=lambda path: common_state_loads.append(path), + ) + import megatron.core + + monkeypatch.setattr(megatron.core, "dist_checkpointing", fake_dist_checkpointing) + args = SimpleNamespace(loss_type="sft", sft_objective="reward_model") + with pytest.raises(RuntimeError, match="RM resume requires a distributed checkpoint"): + checkpoint_module._validate_checkpoint_contract(args, [SimpleNamespace(role="actor")], tmp_path) + assert common_state_loads == [] + + +def test_checkpoint_wrapper_delegates_non_reward_model_legacy_checkpoint(monkeypatch, tmp_path): + import megatron.core + + upstream_loads = [] + fake_dist_checkpointing = SimpleNamespace( + check_is_distributed_checkpoint=lambda path: False, + load_common_state_dict=lambda path: pytest.fail("legacy checkpoint must not use distributed common state"), + ) + args = SimpleNamespace(load=str(tmp_path), loss_type="sft", sft_objective="causal_lm", ckpt_step=None) + monkeypatch.setattr(megatron.core, "dist_checkpointing", fake_dist_checkpointing) + monkeypatch.setattr(checkpoint_module, "get_args", lambda: args) + monkeypatch.setattr(checkpoint_module, "_is_dir_nonempty", lambda _: True) + monkeypatch.setattr(checkpoint_module, "is_megatron_checkpoint", lambda _: True) + monkeypatch.setattr(checkpoint_module, "_checkpoint_iteration_dir", lambda *_: tmp_path / "iter_0000001") + monkeypatch.setattr( + checkpoint_module, + "_load_checkpoint_megatron", + lambda **kwargs: upstream_loads.append(kwargs) or (1, 0), + ) + + assert checkpoint_module.load_checkpoint([SimpleNamespace(role="actor")], None, None, None, False) == (1, 0) + assert len(upstream_loads) == 1 + + +def test_reward_model_contract_accepts_complete_metadata_and_exact_head(monkeypatch, tmp_path): + fake_dist_checkpointing = SimpleNamespace( + check_is_distributed_checkpoint=lambda path: True, + load_common_state_dict=lambda path: { + "args": SimpleNamespace( + sft_objective="reward_model", + head_type="reward_model_terminal_v1", + checkpoint_role="actor", + ) + }, + load_tensors_metadata=lambda path: {"model.output_layer.weight": _TensorMetadata((1, 1024))}, + ) + import megatron.core + + monkeypatch.setattr(megatron.core, "dist_checkpointing", fake_dist_checkpointing) + args = SimpleNamespace( + loss_type="sft", + sft_objective="reward_model", + hidden_size=1024, + no_load_optim=False, + no_load_rng=False, + finetune=False, + reset_optimizer_states=False, + ) + checkpoint_module._validate_checkpoint_contract(args, [SimpleNamespace(role="actor")], tmp_path) + + +@pytest.mark.parametrize("flag", ["no_load_optim", "no_load_rng", "finetune", "reset_optimizer_states"]) +def test_reward_model_contract_rejects_partial_resume_flags(monkeypatch, tmp_path, flag): + fake_dist_checkpointing = SimpleNamespace( + check_is_distributed_checkpoint=lambda path: True, + load_common_state_dict=lambda path: { + "args": SimpleNamespace( + sft_objective="reward_model", + head_type="reward_model_terminal_v1", + checkpoint_role="actor", + ) + }, + ) + import megatron.core + + monkeypatch.setattr(megatron.core, "dist_checkpointing", fake_dist_checkpointing) + values = dict(no_load_optim=False, no_load_rng=False, finetune=False, reset_optimizer_states=False) + values[flag] = True + args = SimpleNamespace(loss_type="sft", sft_objective="reward_model", hidden_size=1024, **values) + with pytest.raises(RuntimeError, match="must restore optimizer, scheduler, and RNG"): + checkpoint_module._validate_checkpoint_contract(args, [SimpleNamespace(role="actor")], tmp_path) + + +def test_restored_scheduler_is_not_advanced_twice(): + args = SimpleNamespace(no_load_optim=False, finetune=False, reset_optimizer_states=False) + assert checkpoint_module.scheduler_state_was_restored(args, resumed_from_megatron=True) + args.no_load_optim = True + assert not checkpoint_module.scheduler_state_was_restored(args, resumed_from_megatron=True) + + +def test_checkpoint_wrapper_restores_optimizer_scheduler_rng_and_next_step_loss(monkeypatch, tmp_path): + torch.manual_seed(7) + source = torch.nn.Linear(3, 1) + source_optimizer = torch.optim.AdamW(source.parameters(), lr=0.01) + source_scheduler = torch.optim.lr_scheduler.StepLR(source_optimizer, step_size=1, gamma=0.5) + + def step(model, optimizer, scheduler): + inputs = torch.randn(4, 3) + targets = torch.randn(4, 1) + optimizer.zero_grad() + loss = torch.nn.functional.mse_loss(model(inputs), targets) + loss.backward() + optimizer.step() + scheduler.step() + return loss.detach() + + step(source, source_optimizer, source_scheduler) + saved_model = copy.deepcopy(source.state_dict()) + saved_optimizer = copy.deepcopy(source_optimizer.state_dict()) + saved_scheduler = copy.deepcopy(source_scheduler.state_dict()) + saved_rng = torch.get_rng_state().clone() + expected_loss = step(source, source_optimizer, source_scheduler) + expected_parameters = copy.deepcopy(source.state_dict()) + + resumed = torch.nn.Linear(3, 1) + resumed_optimizer = torch.optim.AdamW(resumed.parameters(), lr=0.01) + resumed_scheduler = torch.optim.lr_scheduler.StepLR(resumed_optimizer, step_size=1, gamma=0.5) + + def fake_load_checkpoint_megatron(*, ddp_model, optimizer, opt_param_scheduler, **_): + ddp_model[0].load_state_dict(saved_model) + optimizer.load_state_dict(saved_optimizer) + opt_param_scheduler.load_state_dict(saved_scheduler) + torch.set_rng_state(saved_rng) + return 1, 0 + + monkeypatch.setattr(checkpoint_module, "get_args", lambda: SimpleNamespace(load=str(tmp_path))) + monkeypatch.setattr(checkpoint_module, "_is_dir_nonempty", lambda _: True) + monkeypatch.setattr(checkpoint_module, "is_megatron_checkpoint", lambda _: True) + monkeypatch.setattr(checkpoint_module, "_checkpoint_iteration_dir", lambda *_: tmp_path) + monkeypatch.setattr(checkpoint_module, "_validate_checkpoint_contract", lambda *_: None) + monkeypatch.setattr(checkpoint_module, "_load_checkpoint_megatron", fake_load_checkpoint_megatron) + + checkpoint_module.load_checkpoint([resumed], resumed_optimizer, resumed_scheduler, None, False) + actual_loss = step(resumed, resumed_optimizer, resumed_scheduler) + + torch.testing.assert_close(actual_loss, expected_loss) + assert resumed_scheduler.state_dict() == source_scheduler.state_dict() + for name, parameter in resumed.state_dict().items(): + torch.testing.assert_close(parameter, expected_parameters[name]) + + +def test_critic_rejects_reward_model_metadata(monkeypatch, tmp_path): + fake_dist_checkpointing = SimpleNamespace( + check_is_distributed_checkpoint=lambda path: True, + load_common_state_dict=lambda path: { + "args": SimpleNamespace( + sft_objective="reward_model", + head_type="reward_model_terminal_v1", + checkpoint_role="actor", + ) + }, + ) + import megatron.core + + monkeypatch.setattr(megatron.core, "dist_checkpointing", fake_dist_checkpointing) + with pytest.raises(RuntimeError, match="PPO critic load rejects"): + checkpoint_module._validate_checkpoint_contract(SimpleNamespace(), [SimpleNamespace(role="critic")], tmp_path) diff --git a/tests/backends/megatron/test_reward_model_loss.py b/tests/backends/megatron/test_reward_model_loss.py new file mode 100644 index 000000000..bbb6fa856 --- /dev/null +++ b/tests/backends/megatron/test_reward_model_loss.py @@ -0,0 +1,120 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Reward-model loss, pooling, and metric contracts.""" + +from argparse import Namespace +from types import SimpleNamespace + +import pytest +import torch + + +try: + from relax.backends.megatron.loss import reward_model_loss_function + from relax.backends.megatron.model import _restore_micro_batch_output_order +except Exception as exc: + pytest.skip(f"Megatron reward-model loss unavailable: {exc}", allow_module_level=True) + + +def test_preference_outputs_restore_original_order_without_dynamic_batch_flag(): + packed_values = [ + "pair-2-chosen", + "pair-2-rejected", + "pair-0-chosen", + "pair-0-rejected", + "pair-1-chosen", + "pair-1-rejected", + ] + schedule = [[4, 5, 0, 1], [2, 3]] + + assert _restore_micro_batch_output_order(packed_values, schedule) == [ + "pair-0-chosen", + "pair-0-rejected", + "pair-1-chosen", + "pair-1-rejected", + "pair-2-chosen", + "pair-2-rejected", + ] + assert _restore_micro_batch_output_order(["aggregate-0", "aggregate-1"], schedule) == [ + "aggregate-0", + "aggregate-1", + ] + + +def _batch(): + branches = [ + torch.tensor([10, 11, 12]), + torch.tensor([20, 21]), + torch.tensor([30, 31, 32, 33]), + torch.tensor([40, 41]), + ] + lengths = [len(branch) for branch in branches] + packed = torch.cat(branches) + return { + "total_lengths": lengths, + "response_lengths": lengths, + "score_positions": [2, 1, 3, 1], + "raw_loss_masks": [ + torch.tensor([0, 1, 1]), + torch.tensor([0, 1]), + torch.tensor([0, 0, 1, 1]), + torch.tensor([0, 1]), + ], + "tokens": packed.unsqueeze(0), + "unconcat_tokens": branches, + "packed_seq_params": SimpleNamespace(cu_seqlens_q=torch.tensor([0, 3, 5, 9, 11])), + "preference_branch_pair_ids": [7, 8, 7, 8], + "preference_is_chosen": [False, True, True, False], + } + + +def test_reward_model_loss_uses_pair_identity_after_branch_reordering_and_preserves_gradient(): + batch = _batch() + flat = torch.arange(11, dtype=torch.float32, requires_grad=True) + loss, metrics = reward_model_loss_function(Namespace(), batch, flat.reshape(1, 11, 1), lambda value: value) + + expected_margins = torch.tensor([8.0 - 2.0, 4.0 - 10.0]) + expected = -torch.nn.functional.logsigmoid(expected_margins) + assert torch.allclose(loss, expected.sum()) + assert set(metrics) == { + "rm/loss", + "rm/score_chosen_mean", + "rm/score_rejected_mean", + "rm/score_margin_mean", + "rm/accuracy", + "rm/_score_chosen_second_moment", + "rm/_score_rejected_second_moment", + } + loss.backward() + assert flat.grad is not None + assert torch.count_nonzero(flat.grad).item() == 4 + + +@pytest.mark.parametrize( + ("mutation", "match"), + [ + (lambda batch: batch["raw_loss_masks"][0].zero_(), "raw_loss_mask"), + (lambda batch: batch["tokens"][0].__setitem__(2, 999), "terminal token"), + ( + lambda batch: setattr(batch["packed_seq_params"], "cu_seqlens_q", torch.tensor([0, 2, 5, 9, 11])), + "cu_seqlens", + ), + ], +) +def test_reward_model_pooling_rejects_mask_token_and_segment_misalignment(mutation, match): + batch = _batch() + mutation(batch) + with pytest.raises(ValueError, match=match): + reward_model_loss_function(Namespace(), batch, torch.zeros(1, 11, 1), lambda value: value) + + +def test_reward_model_pooling_allows_only_one_trailing_padding_segment(): + batch = _batch() + batch["tokens"] = torch.nn.functional.pad(batch["tokens"], (0, 5)) + batch["packed_seq_params"].cu_seqlens_q = torch.tensor([0, 3, 5, 9, 11, 16]) + logits = torch.zeros(1, 16, 1) + reward_model_loss_function(Namespace(), batch, logits, lambda value: value) + + batch["packed_seq_params"].cu_seqlens_q = torch.tensor([0, 3, 5, 9, 11, 14, 16]) + with pytest.raises(ValueError, match="at most one trailing padding"): + reward_model_loss_function(Namespace(), batch, logits, lambda value: value) diff --git a/tests/backends/megatron/test_sft_train_actor_eval.py b/tests/backends/megatron/test_sft_train_actor_eval.py index f3869e74b..69752b96e 100644 --- a/tests/backends/megatron/test_sft_train_actor_eval.py +++ b/tests/backends/megatron/test_sft_train_actor_eval.py @@ -5,21 +5,13 @@ from argparse import Namespace -import pytest - - -# Importing relax.backends.megatron.actor pulls in CUDA-only deps. Skip the -# whole module on CPU-only envs — matches the pattern used in -# tests/backends/megatron/test_sft_train_data_fields.py. -try: - from relax.backends.megatron.actor import _should_run_sft_eval # noqa: F401 -except (ImportError, AssertionError) as _exc: - pytest.skip(f"relax.backends.megatron.actor unavailable: {_exc}", allow_module_level=True) +from relax.engine.sft.runtime import evaluation_step_for_rollout, should_run_sft_eval, should_run_sft_predict def _mk_actor_args(): return Namespace( loss_type="sft", + sft_objective="causal_lm", compute_advantages_and_returns=False, eval_prompt_data=["eval", "/dev/null"], eval_size=None, @@ -36,25 +28,55 @@ def _mk_actor_args(): def test_should_run_sft_eval_at_interval_boundary(): args = _mk_actor_args() - assert _should_run_sft_eval(args, rollout_id=9) is True - assert _should_run_sft_eval(args, rollout_id=19) is True - assert _should_run_sft_eval(args, rollout_id=4) is False - assert _should_run_sft_eval(args, rollout_id=0) is False + assert should_run_sft_eval(args, completed_steps=10) is True + assert should_run_sft_eval(args, completed_steps=20) is True + assert should_run_sft_eval(args, completed_steps=5) is False + assert should_run_sft_eval(args, completed_steps=0) is False + + +def test_preference_eval_includes_true_baseline_and_final_independent_of_interval(): + args = _mk_actor_args() + args.sft_objective = "reward_model" + args.eval_interval = 200 + assert should_run_sft_eval(args, completed_steps=0) is True + assert should_run_sft_eval(args, completed_steps=20) is True + assert should_run_sft_eval(args, completed_steps=1) is False def test_should_run_sft_eval_disabled_when_no_interval(): args = _mk_actor_args() args.eval_interval = None - assert _should_run_sft_eval(args, rollout_id=9) is False + assert should_run_sft_eval(args, completed_steps=10) is False def test_should_run_sft_eval_disabled_when_no_eval_source(): args = _mk_actor_args() args.eval_prompt_data = None - assert _should_run_sft_eval(args, rollout_id=9) is False + assert should_run_sft_eval(args, completed_steps=10) is False def test_should_run_sft_eval_disabled_for_non_sft(): args = _mk_actor_args() args.loss_type = "policy_loss" - assert _should_run_sft_eval(args, rollout_id=9) is False + assert should_run_sft_eval(args, completed_steps=10) is False + + +def test_should_run_sft_predict_uses_completed_step_boundaries(): + args = _mk_actor_args() + args.sft_predict_interval = 10 + + assert should_run_sft_predict(args, completed_steps=0) is False + assert should_run_sft_predict(args, completed_steps=9) is False + assert should_run_sft_predict(args, completed_steps=10) is True + assert should_run_sft_predict(args, completed_steps=19) is False + assert should_run_sft_predict(args, completed_steps=20) is True + + +def test_evaluation_step_mapping_is_completed_for_sft_and_zero_based_for_rl(): + args = _mk_actor_args() + assert evaluation_step_for_rollout(args, rollout_id=0) == 1 + assert evaluation_step_for_rollout(args, rollout_id=9) == 10 + + args.loss_type = "policy_loss" + assert evaluation_step_for_rollout(args, rollout_id=0) == 0 + assert evaluation_step_for_rollout(args, rollout_id=9) == 9 diff --git a/tests/backends/megatron/test_sft_train_data_fields.py b/tests/backends/megatron/test_sft_train_data_fields.py index 21a58d111..579075f21 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,55 @@ 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", + "chosen_score_positions", + "rejected_score_positions", + ] + + +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], + "chosen_score_positions": [2], + "rejected_score_positions": [1], + } + 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 "pair_ids" not in rollout_data + assert rollout_data["preference_pair_ids"] == [17] + assert rollout_data["preference_branch_pair_ids"] == [17, 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/components/test_actor_sft_partition.py b/tests/components/test_actor_sft_partition.py index 19a286e03..88225169e 100644 --- a/tests/components/test_actor_sft_partition.py +++ b/tests/components/test_actor_sft_partition.py @@ -44,6 +44,33 @@ def test_actor_helpers_emit_train_partition_under_rl(): assert sft_task_name(rl_cfg, component="actor") == "train_actor" +def test_colocate_preference_step_zero_starts_when_baseline_partition_is_ready(): + from relax.engine.sft.runtime import actor_training_input_ready + + cfg = _mk_actor_config(loss_type="sft") + cfg.sft_objective = "reward_model" + cfg.eval_prompt_data = ["task31", "heldout.parquet"] + cfg.eval_size = None + cfg.eval_interval = 200 + + assert actor_training_input_ready(cfg, 0, ["sft_eval_0_n2_0"]) + assert not actor_training_input_ready(cfg, 0, []) + assert not actor_training_input_ready(cfg, 0, ["sft_0"]) + assert not actor_training_input_ready(cfg, 0, ["sft_eval_0_n2_1"]) + assert not actor_training_input_ready(cfg, 1, ["sft_eval_1_n2_0"]) + assert actor_training_input_ready(cfg, 1, ["sft_1"]) + + +def test_actor_training_partition_remains_the_normal_readiness_signal(): + from relax.engine.sft.runtime import actor_training_input_ready + + cfg = _mk_actor_config(loss_type="sft") + cfg.sft_objective = "causal_lm" + + assert actor_training_input_ready(cfg, 0, ["sft_0"]) + assert not actor_training_input_ready(cfg, 0, ["sft_eval_0_n2_0"]) + + def test_actor_resume_uses_backend_step_over_auto_cold_start(): from relax.components.actor import _resolve_start_rollout_id diff --git a/tests/components/test_sft.py b/tests/components/test_sft.py index 213a654e6..322b536d0 100644 --- a/tests/components/test_sft.py +++ b/tests/components/test_sft.py @@ -178,6 +178,7 @@ async def test_sft_eval_rejects_source_with_no_valid_samples(monkeypatch): args = _make_args(global_batch_size=4) args.eval_interval = 1 + args.eval_prompt_data = "/fake/eval.jsonl" SFTCls = SFT.func_or_class sft = SFTCls.__new__(SFTCls) sft.config = args @@ -194,7 +195,7 @@ async def test_sft_eval_rejects_source_with_no_valid_samples(monkeypatch): sft._stop_event.is_set = MagicMock(return_value=False) with pytest.raises(RuntimeError, match="source produced 0 valid samples"): - await sft._maybe_produce_eval() + await sft._maybe_produce_eval(completed_steps=1) fake_client.async_put.assert_not_awaited() 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..42454a3f7 --- /dev/null +++ b/tests/engine/sft/dataset/test_preference.py @@ -0,0 +1,225 @@ +# 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() + assert pair.chosen_score_position == pair.chosen_total_length - 1 + assert pair.rejected_score_position == pair.rejected_total_length - 1 + + +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/eval/test_acceptance.py b/tests/engine/sft/eval/test_acceptance.py new file mode 100644 index 000000000..906550282 --- /dev/null +++ b/tests/engine/sft/eval/test_acceptance.py @@ -0,0 +1,107 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Golden tests for RFC paired-bootstrap and artifact semantics.""" + +import json +from types import SimpleNamespace + +import numpy as np +import pytest + +from relax.engine.sft.eval.acceptance import ( + encoded_pair_id, + paired_bootstrap, + preference_eval_chunk_sizes, + preference_eval_local_batch_sizes, + record_probe_contract, + write_pair_artifacts, +) + + +@pytest.mark.parametrize( + ("global_batch_size", "expected"), + [ + (1, [1] * 512), + (30, [30] * 17 + [2]), + (32, [32] * 16), + (512, [512]), + (513, [512]), + ], +) +def test_preference_eval_chunks_preserve_all_512_unique_pairs(global_batch_size, expected): + sizes = preference_eval_chunk_sizes(512, global_batch_size) + assert sizes == expected + assert sum(sizes) == 512 + assert max(sizes) <= global_batch_size + + +def test_preference_eval_partial_chunk_uses_actual_per_rank_batch_size(): + assert preference_eval_local_batch_sizes(512, 30, dp_size=2) == [15] * 17 + [1] + with pytest.raises(ValueError, match="divisible by data-parallel size"): + preference_eval_local_batch_sizes(512, 31, dp_size=2) + + +def test_paired_bootstrap_matches_pcg64_float64_golden_fixture(): + result = paired_bootstrap([1.0, 0.0, 1.0, 0.5]) + assert result["point_estimate"] == 0.625 + assert result["lower_95"] == 0.25 + assert result["indices_sha256"] == "dc7ec9501aead17b5115aad49b487302aed95c384d6c2aadcf67b56a849ef53f" + assert result["replicates_sha256"] == "52f63cf64c72be2904c2911052092af8fc2cd2f05444fa17898c79a1dc22e174" + assert result["passes_lower_bound_gt_0_50"] is False + + +def _pair(pair_id: str, token: int): + return SimpleNamespace( + pair_id=pair_id, + chosen_tokens=np.asarray([1, token]), + rejected_tokens=np.asarray([1, token + 1]), + chosen_loss_mask=np.asarray([0, 1]), + rejected_loss_mask=np.asarray([0, 1]), + chosen_score_position=1, + rejected_score_position=1, + ) + + +def test_pair_artifacts_preserve_original_ids_and_require_identical_final_plan(tmp_path): + pairs = [_pair(f"pair-{index}", index) for index in range(4)] + record_probe_contract(str(tmp_path), "reward_model", 0, pairs, expected_pair_count=4) + rows = [ + { + "encoded_pair_id": encoded_pair_id(pair.pair_id), + "chosen_score": float(index + 1), + "rejected_score": 0.0, + "pair_loss": 0.1, + } + for index, pair in enumerate(pairs) + ] + plan = [ + {"rank": 0, "chunk": 0, "batch": 0, "microbatch": 0, "encoded_pair_ids": [r["encoded_pair_id"] for r in rows]} + ] + write_pair_artifacts(str(tmp_path), "reward_model", 0, rows, plan) + record_probe_contract(str(tmp_path), "reward_model", 4, pairs, expected_pair_count=4) + write_pair_artifacts(str(tmp_path), "reward_model", 4, rows, plan) + + pair_path = tmp_path / "preference_eval" / "reward_model-step-0000004-pairs.jsonl" + written = [json.loads(line) for line in pair_path.read_text(encoding="utf-8").splitlines()] + assert [row["pair_id"] for row in written] == [pair.pair_id for pair in pairs] + assert set(written[0]) == {"pair_id", "chosen_score", "rejected_score", "pair_loss"} + + changed_plan = [{**plan[0], "encoded_pair_ids": list(reversed(plan[0]["encoded_pair_ids"]))}] + with pytest.raises(RuntimeError, match="batch plan changed"): + write_pair_artifacts(str(tmp_path), "reward_model", 4, rows, changed_plan) + + +def test_probe_contract_rejects_any_count_other_than_frozen_512(tmp_path): + with pytest.raises(RuntimeError, match="exactly 512 probe pairs"): + record_probe_contract(str(tmp_path), "reward_model", 0, [_pair("pair-0", 0)]) + + +@pytest.mark.parametrize("save_path", [None, ""]) +def test_probe_contract_rejects_nonfrozen_count_without_artifact_output(save_path): + with pytest.raises(RuntimeError, match="exactly 512 probe pairs"): + record_probe_contract(save_path, "reward_model", 0, [_pair("pair-0", 0)]) + + +def test_probe_contract_accepts_frozen_count_without_artifact_output(): + pairs = [_pair(f"pair-{index}", index) for index in range(512)] + assert record_probe_contract(None, "reward_model", 0, pairs) is None diff --git a/tests/engine/sft/eval/test_preference.py b/tests/engine/sft/eval/test_preference.py new file mode 100644 index 000000000..4dbcc6fda --- /dev/null +++ b/tests/engine/sft/eval/test_preference.py @@ -0,0 +1,54 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +import pytest +import torch + +from relax.engine.sft.eval.preference import ( + compute_reward_model_eval_step, + extract_preference_eval_pair_ids, + finalize_pair_metrics, + pair_metric_sums, +) + + +@pytest.mark.parametrize( + ("pair_data", "expected"), + [ + ({"pair_ids": [17, 23]}, [17, 23]), + ({"preference_pair_ids": [17, 23]}, [17, 23]), + ({"pair_ids": [99], "preference_pair_ids": [17]}, [17]), + ], +) +def test_extract_preference_eval_pair_ids_supports_raw_and_expanded_data(pair_data, expected): + assert extract_preference_eval_pair_ids(pair_data) == expected + + +def test_extract_preference_eval_pair_ids_rejects_missing_pair_level_identity(): + with pytest.raises(RuntimeError, match="expected preference_pair_ids or pair_ids"): + extract_preference_eval_pair_ids({"preference_branch_pair_ids": [17, 17]}) + + +def test_reward_model_eval_emits_one_score_per_branch_for_order_restoration(): + _, outputs = compute_reward_model_eval_step( + torch.tensor([0.0, 1.0, 2.0, 3.0]), + total_lengths=[2, 2], + score_positions=[1, 1], + ) + + assert len(outputs["scores"]) == 2 + assert all(score.ndim == 0 for score in outputs["scores"]) + assert torch.stack(outputs["scores"]).tolist() == [1.0, 3.0] + + +def test_pair_metric_sums_and_finalize_keep_ties_explicit(): + chosen = torch.tensor([2.0, 1.0, 1.0]) + rejected = torch.tensor([1.0, 2.0, 1.0]) + losses = torch.tensor([0.1, 0.2, 0.3]) + + metrics = finalize_pair_metrics(pair_metric_sums(chosen, rejected, losses), prefix="rm") + + assert metrics["eval/rm_loss"] == pytest.approx(0.2) + assert metrics["eval/rm_strict_accuracy"] == pytest.approx(1 / 3) + assert metrics["eval/rm_tie_rate"] == pytest.approx(1 / 3) + assert metrics["eval/rm_tie_aware_accuracy"] == pytest.approx(0.5) + assert metrics["eval/rm_pairs"] == 3 diff --git a/tests/engine/sft/test_preference_runtime.py b/tests/engine/sft/test_preference_runtime.py new file mode 100644 index 000000000..679c9e9eb --- /dev/null +++ b/tests/engine/sft/test_preference_runtime.py @@ -0,0 +1,115 @@ +# 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, + "save_hf": None, + "eval_interval": 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")) + + +@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"), + ], +) +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)) + + +def test_reward_model_accepts_hf_load_and_held_out_evaluation(): + validate_preference_args( + _args( + sft_objective="reward_model", + ref_load="/models/Qwen3-0.6B", + eval_prompt_data=["task31", "heldout.parquet"], + rollout_temperature=0.8, + enable_weights_backuper=False, + dpo_reference_repository=None, + dpo_reference_revision=None, + save_hf=None, + ) + ) + + +def test_reward_model_rejects_hf_export(): + with pytest.raises(ValueError, match="reward_model v1 does not support --save-hf"): + validate_preference_args(_args(sft_objective="reward_model", save_hf="/models/rm-{rollout_id}")) + + +def test_dpo_and_causal_sft_keep_hf_export_support(): + validate_preference_args(_args(save_hf="/models/dpo-{rollout_id}")) + validate_preference_args(_args(sft_objective="causal_lm", save_hf="/models/sft-{rollout_id}")) diff --git a/tests/utils/training/test_preference_utils.py b/tests/utils/training/test_preference_utils.py new file mode 100644 index 000000000..5c91308cf --- /dev/null +++ b/tests/utils/training/test_preference_utils.py @@ -0,0 +1,152 @@ +# 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, + reward_model_pair_loss, + select_packed_sequence_scores, +) + + +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_reward_model_pair_loss_matches_independent_reference_and_gradient(): + chosen = torch.tensor([1.0, -1.0, 0.0], requires_grad=True) + rejected = torch.tensor([0.0, 2.0, 0.0], requires_grad=True) + + actual = reward_model_pair_loss(chosen, rejected) + expected = -F.logsigmoid(chosen - rejected) + assert torch.allclose(actual, expected, rtol=1e-6, atol=1e-6) + + actual.sum().backward() + actual_grad = chosen.grad.detach().clone() + chosen.grad = None + expected.sum().backward() + assert torch.allclose(actual_grad, chosen.grad, rtol=1e-6, atol=1e-6) + + +def test_reward_model_pair_loss_uses_shared_tensor_condition(monkeypatch): + calls = [] + monkeypatch.setattr( + "relax.utils.training.preference_utils.require_tensor_condition", + lambda condition, message: calls.append((condition, message)), + ) + + reward_model_pair_loss(torch.tensor([1.0]), torch.tensor([0.0])) + + assert len(calls) == 1 + assert "finite" in calls[0][1] + + +def test_reward_model_pair_loss_rejects_non_finite_cpu_margin(): + with pytest.raises(ValueError, match="finite"): + reward_model_pair_loss(torch.tensor([float("inf")]), torch.tensor([0.0])) + + +def test_select_packed_sequence_scores_preserves_pair_order_and_gradient(): + flat_logits = torch.arange(10, dtype=torch.float32, requires_grad=True) + logits = flat_logits.reshape(1, 10, 1) + + scores = select_packed_sequence_scores(logits, [3, 2, 4], [2, 1, 3]) + + assert scores.tolist() == [2.0, 4.0, 8.0] + scores.sum().backward() + expected = torch.zeros(10) + expected[[2, 4, 8]] = 1 + assert torch.equal(flat_logits.grad, expected) + + +def test_select_packed_sequence_scores_rejects_invalid_position(): + with pytest.raises(ValueError, match="outside"): + select_packed_sequence_scores(torch.zeros(1, 4, 1), [4], [4]) + + +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)